From 537822d1b8b5fcc899cab3010c320159c21d1bbd Mon Sep 17 00:00:00 2001 From: Zack Meeks Date: Mon, 17 Aug 2026 22:00:34 +0000 Subject: [PATCH 01/21] Fix GPU and Java-heap OOM when building large vector indexes Signed-off-by: Zack Meeks --- .../cuvs/lucene/AcceleratedHNSWUtils.java | 54 ++++++---------- .../cuvs/lucene/CuVS2510GPUVectorsWriter.java | 6 +- .../Lucene99AcceleratedHNSWVectorsWriter.java | 62 ++++++++++++++----- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 6 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 6 +- .../java/com/nvidia/cuvs/lucene/Utils.java | 62 +++++-------------- 6 files changed, 89 insertions(+), 107 deletions(-) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 9c49c07fe0..1be9f2450c 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -77,27 +77,32 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens * M = ceil(cagraGraphDegree / 2), where cagraGraphDegree is the CAGRA adjacency list's degree * (its column count). Ceil is used to accommodate odd graph degrees. * Each layer contains 1/M nodes from the previous layer - * Creates layers until the highest layer has ≤ M nodes + * Creates layers until the highest layer has <= M nodes + *

+ * Vectors for higher-layer subsets are read directly from the native matrix + * via {@link CuVSMatrix#getRow(long)} and {@link RowView#toArray(float[])}, + * avoiding any additional heap allocation of the full dataset. Used by both + * the flush and merge paths; the caller provides the vectors as a + * {@link CuVSMatrix}. */ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, - int size, int dimensions, CuVSMatrix adjacencyListMatrix, - List vectors, + CuVSMatrix vectorDataset, int hnswLayers, CagraIndexParams params, QuantizationType quantization) throws Throwable { + int size = (int) vectorDataset.size(); int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); - // Store all layers data List layerNodes = new ArrayList<>(); List layerAdjacencies = new ArrayList<>(); // Layer 0: Use full CAGRA adjacency list - layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list + layerNodes.add(null); layerAdjacencies.add(adjacencyListMatrix); int currentLayerSize = size; @@ -105,67 +110,50 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( Random random = new Random(); while (layerIndex < hnswLayers && currentLayerSize > 1) { - // Calculate size for next layer (1/M of current layer) int nextLayerSize = Math.max(2, currentLayerSize / M); - // Select nodes for this layer SortedSet selectedNodesSet = new TreeSet<>(); if (layerIndex == 1) { - // Select from all nodes (Layer 0) while (selectedNodesSet.size() < nextLayerSize) { selectedNodesSet.add(random.nextInt(size)); } } else { - // Select from previous layer nodes int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); while (selectedNodesSet.size() < nextLayerSize) { - int idx = random.nextInt(prevLayerNodes.length); - selectedNodesSet.add(prevLayerNodes[idx]); + selectedNodesSet.add(prevLayerNodes[random.nextInt(prevLayerNodes.length)]); } } - // Convert to sorted array int[] selectedNodes = selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray(); - layerNodes.add(selectedNodes); if (quantization == QuantizationType.NONE) { - // Extract vectors for selected nodes - float[][] selectedVectors = new float[nextLayerSize][]; + // Read only the sampled rows from the native matrix — no full-dataset heap copy + float[][] selectedVectors = new float[nextLayerSize][dimensions]; for (int i = 0; i < nextLayerSize; i++) { - selectedVectors[i] = (float[]) vectors.get(selectedNodes[i]); + vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); } - - // Build CAGRA graph for this layer layerAdjacencies.add( buildCagraGraphForSubset( selectedVectors, selectedNodes, 0, params, dimensions, quantization)); - } else { - - // Extract vectors for selected nodes - int bytesPerVector = (dimensions + 7) / 8; - byte[][] selectedVectors = new byte[nextLayerSize][]; + // Byte width comes from the matrix itself: binary packs 8 dims/byte, scalar is 1 byte/dim. + int bytesPerVector = (int) vectorDataset.columns(); + byte[][] selectedVectors = new byte[nextLayerSize][bytesPerVector]; for (int i = 0; i < nextLayerSize; i++) { - selectedVectors[i] = (byte[]) vectors.get(selectedNodes[i]); + vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); } - - // Build CAGRA graph for this layer layerAdjacencies.add( buildCagraGraphForSubset( selectedVectors, selectedNodes, bytesPerVector, params, dimensions, quantization)); } - // Update for next iteration currentLayerSize = nextLayerSize; layerIndex++; - - // Use different seed for each layer random = new Random(new Random().nextLong()); } - // Create the multi-layer graph with all layers return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); } @@ -184,11 +172,9 @@ private static CuVSMatrix buildCagraGraphForSubset( CuVSMatrix subsetDataset; if (quantization == QuantizationType.BINARY) { - subsetDataset = - createByteMatrixFromArray((byte[][]) vectors, bytesPerVector, getCuVSResourcesInstance()); + subsetDataset = createByteMatrixFromArray((byte[][]) vectors, bytesPerVector); } else if (quantization == QuantizationType.SCALAR) { - subsetDataset = - createByteMatrixFromArray((byte[][]) vectors, dimensions, getCuVSResourcesInstance()); + subsetDataset = createByteMatrixFromArray((byte[][]) vectors, dimensions); } else { subsetDataset = CuVSMatrix.ofArray((float[][]) vectors); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 97fedf924f..ce77f36dea 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -201,8 +201,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); try { CuVSMatrix cagraDataset = - Utils.createFloatMatrix( - vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); writeCagraIndex(cagraIndexOutputStream, cagraDataset); } catch (Throwable t) { // Fallback to brute force in a few cases, for now. @@ -215,8 +214,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro if (indexType.isBruteForce()) { var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); CuVSMatrix bruteforceDataset = - Utils.createFloatMatrix( - vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); writeBruteForceIndex(bruteForceIndexOutputStream, bruteforceDataset); bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 13edc64975..e3e10f44a1 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -16,12 +16,12 @@ import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; -import static com.nvidia.cuvs.lucene.Utils.createListFromMergedVectors; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; import java.io.IOException; @@ -34,11 +34,14 @@ import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; import org.apache.lucene.index.DocsWithFieldSet; 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.Sorter.DocMap; +import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.IOUtils; import org.apache.lucene.util.InfoStream; @@ -139,7 +142,8 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException } /** - * Builds the intermediate CAGRA index and builds and writes the HNSW index. + * Flush/sorting path: builds a host matrix from the heap vectors, then delegates + * to {@link #writeFieldInternal(FieldInfo, CuVSMatrix)}. * * @param fieldInfo instance of FieldInfo that has the field description * @param vectors vectors to index @@ -154,29 +158,48 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro writeSingleVectorGraph(fieldInfo, vectors); return; } - try { - CuVSMatrix dataset = - Utils.createFloatMatrix( - vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeFieldInternal(fieldInfo, dataset); + } + /** + * Builds the intermediate CAGRA index and builds and writes the HNSW index. + * Single implementation used by both the flush and merge paths. The dataset is a + * {@link CuVSMatrix} (host-backed on the merge path) so the full set of vectors is + * never double-materialised on the Java heap. + * + * @param fieldInfo instance of FieldInfo that has the field description + * @param dataset matrix of all vectors to index + * @throws IOException + */ + private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { + int size = (int) dataset.size(); + if (size == 0) { + writeEmpty(fieldInfo, hnswMeta); + return; + } + if (size < 2) { + float[] buf = new float[fieldInfo.getVectorDimension()]; + dataset.getRow(0).toArray(buf); + writeSingleVectorGraph(fieldInfo, List.of(buf)); + return; + } + try { CagraIndexParams params = CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); - CagraIndex cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()) .withDataset(dataset) .withIndexParams(params) .build(); CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); - int size = (int) dataset.size(); int dimensions = fieldInfo.getVectorDimension(); GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, - size, dimensions, adjacencyListMatrix, - vectors, + dataset, acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.NONE); @@ -273,13 +296,24 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) } /** - * Create combined data set for the merged segment and call writeFieldInternal. + * Streams merged vectors directly into a native host-memory matrix (CuVSHostMatrix) + * without materialising a List on the Java heap, then calls writeFieldInternal. + * This avoids the double-copy OOM (heap list + native matrix simultaneously) that + * occurs when force-merging large segments. */ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { try { - List dataset = - createListFromMergedVectors( - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState)); + FloatVectorValues mergedVectors = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + int size = mergedVectors.size(); + int dims = fieldInfo.getVectorDimension(); + CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT); + KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); + for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { + builder.addVector(mergedVectors.vectorValue(it.index())); + } + CuVSHostMatrix dataset = builder.build(); writeFieldInternal(fieldInfo, dataset); } catch (Throwable t) { Utils.handleThrowable(t); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 87907d2cbb..10380ca538 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -155,8 +155,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw int dimensions = fieldInfo.getVectorDimension(); int bytesPerVector = (dimensions + 7) / 8; - CuVSMatrix dataset = - Utils.createByteMatrix(vectors, bytesPerVector, getCuVSResourcesInstance()); + CuVSMatrix dataset = Utils.createByteMatrix(vectors, bytesPerVector); if (dataset.size() < 2) { writeSingleVectorGraph(fieldInfo, vectors); @@ -179,10 +178,9 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, - size, dimensions, adjacencyListMatrix, - vectors, + dataset, acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.BINARY); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 7141af56ee..f9dbbfb06d 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -181,8 +181,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE } // Create CuVSMatrix with BYTE data type (unsigned bytes) - CuVSMatrix dataset = - Utils.createByteMatrix(unsignedVectors, dimensions, getCuVSResourcesInstance()); + CuVSMatrix dataset = Utils.createByteMatrix(unsignedVectors, dimensions); if (dataset.size() < 2) { writeSingleVectorGraph(fieldInfo, unsignedVectors); @@ -204,10 +203,9 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( fieldInfo, - size, dimensions, adjacencyListMatrix, - unsignedVectors, + dataset, acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.SCALAR); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java index e4a20d2b4d..0fa96ded3e 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -43,82 +43,50 @@ static void handleThrowable(Throwable t) throws IOException { } /** - * A method to build a CuVSMatrix from a list of float vectors. + * Builds a host-memory CuVSMatrix from a list of float vectors. * - * Uses CuVSMatrix.Builder to copy vectors directly to device memory - * without creating intermediate heap arrays. + *

Copies vectors directly into a native host matrix via {@link CuVSMatrix#hostBuilder}, + * without creating an intermediate {@code float[][]} on the heap. * * @param data The float vectors - * @param dimensions The number float elements in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix + * @param dimensions The number of float elements in each vector + * @return a host-memory CuVSMatrix */ - static CuVSMatrix createFloatMatrix(List data, int dimensions, CuVSResources resources) { - // Use Builder pattern to avoid intermediate float[][] allocation - // and copy directly from List to device memory + static CuVSMatrix createFloatMatrix(List data, int dimensions) { CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.size(), // rows (number of vectors) - dimensions, // columns (vector dimension) - CuVSMatrix.DataType.FLOAT); - - // Add vectors one by one - builder copies directly to device memory + CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT); for (float[] vector : data) { builder.addVector(vector); } - return builder.build(); } /** - * A method to build a CuVSMatrix from a list of byte vectors (for binary quantized vectors). - * - * Uses CuVSMatrix.Builder to copy vectors directly to device memory - * without creating intermediate heap arrays. + * Builds a host-memory CuVSMatrix from a list of byte vectors (e.g. quantized vectors). * * @param data The byte vectors (packed bits for binary quantization) * @param bytesPerVector The number of bytes in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix with BYTE data type + * @return a host-memory CuVSMatrix with BYTE data type */ - static CuVSMatrix createByteMatrix( - List data, int bytesPerVector, CuVSResources resources) { - // Use Builder pattern to avoid intermediate byte[][] allocation - // and copy directly from List to device memory + static CuVSMatrix createByteMatrix(List data, int bytesPerVector) { CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.size(), // rows (number of vectors) - bytesPerVector, // columns (bytes per vector) - CuVSMatrix.DataType.BYTE); - - // Add vectors one by one - builder copies directly to device memory + CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE); for (byte[] vector : data) { builder.addVector(vector); } - return builder.build(); } /** - * A method to build a CuVSMatrix from a 2D byte array (for binary quantized vectors). + * Builds a host-memory CuVSMatrix from a 2D byte array (e.g. quantized vectors). * * @param data The 2D byte array (packed bits for binary quantization) * @param bytesPerVector The number of bytes in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix with BYTE data type + * @return a host-memory CuVSMatrix with BYTE data type */ - static CuVSMatrix createByteMatrixFromArray( - byte[][] data, int bytesPerVector, CuVSResources resources) { + static CuVSMatrix createByteMatrixFromArray(byte[][] data, int bytesPerVector) { CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.length, // rows (number of vectors) - bytesPerVector, // columns (bytes per vector) - CuVSMatrix.DataType.BYTE); - - // Add vectors one by one - builder copies directly to device memory + CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE); for (byte[] vector : data) { builder.addVector(vector); } From b10b7bdbff738ef620d0231c77bf1f44f87b4617 Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Wed, 16 Sep 2026 08:40:48 +0000 Subject: [PATCH 02/21] Parallelize bounded CAGRA graph materialization --- .../cuvs/lucene/AcceleratedHNSWUtils.java | 7 +- .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 118 ++++++++- .../Lucene99AcceleratedHNSWVectorsWriter.java | 3 +- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 3 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 3 +- ...TestWriterThreadsGraphMaterialization.java | 232 ++++++++++++++++++ 6 files changed, 350 insertions(+), 16 deletions(-) create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 1be9f2450c..d5dfebb486 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -69,7 +69,7 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens layerAdjacencies.add(adjacencyMatrix); // Create the single-layer graph - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, 1); } /** @@ -92,7 +92,8 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( CuVSMatrix vectorDataset, int hnswLayers, CagraIndexParams params, - QuantizationType quantization) + QuantizationType quantization, + int numThreads) throws Throwable { int size = (int) vectorDataset.size(); @@ -154,7 +155,7 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( random = new Random(new Random().nextLong()); } - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, numThreads); } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index 7e9f888e32..22e47b0f83 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -6,10 +6,17 @@ import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.RowView; +import java.io.IOException; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import org.apache.lucene.search.TaskExecutor; import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.NeighborArray; @@ -38,9 +45,15 @@ public class GPUBuiltHnswGraph extends HnswGraph { * @param dimensions the vector dimension * @param layerNodes the nodes on the layer * @param layerAdjacencies adjacency list + * @param numThreads threads to use for materializing the adjacency (1 = serial) */ public GPUBuiltHnswGraph( - int size, int dimensions, List layerNodes, List layerAdjacencies) { + int size, + int dimensions, + List layerNodes, + List layerAdjacencies, + int numThreads) + throws IOException { this.size = size; this.dimensions = dimensions; @@ -50,38 +63,123 @@ public GPUBuiltHnswGraph( // Process Layer 0 (base layer with all nodes) CuVSMatrix layer0Adjacency = layerAdjacencies.get(0); - this.layer0Neighbors = fillNeighborArray(layer0Adjacency, size); + this.layer0Neighbors = fillNeighborArray(layer0Adjacency, size, numThreads); // Process higher layers (1 to numLevels-1) for (int level = 1; level < numLevels; level++) { int[] nodes = layerNodes.get(level); CuVSMatrix adjacency = layerAdjacencies.get(level); this.layerNodes.add(nodes); - this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length)); + this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length, numThreads)); } } + /** Node count below which parallel materialization is not worth the thread overhead. */ + static final int PARALLEL_MIN_NODES = 1 << 16; + + /** + * Maximum temporary native-host copy used to make a device adjacency safe for concurrent reads. + * Larger device matrices retain serial row access instead of risking a full-matrix native-memory + * spike on top of the Java {@link NeighborArray} representation. + */ + static final long MAX_PARALLEL_GRAPH_COPY_BYTES = 4L << 30; + /** - * Fills the neighbor array using the adjacency matrix. + * Materializes the adjacency matrix into on-heap {@link NeighborArray}s, one per node. + * + *

The serial path reads the adjacency directly (a device matrix's {@code getRow} is safe + * single-threaded). The parallel path cannot: the CAGRA layer-0 adjacency is a device matrix whose + * {@code getRow} uses a shared, stateful buffered reader that is not safe for concurrent access, so + * it is pulled to host once (a single bulk device->host copy) before materializing disjoint node + * ranges concurrently. Host matrices (the upper layers, built via {@link CuVSMatrix#ofArray}) are + * read directly in both paths. * * @param adjacency instance of adjacency CuVSMatrix * @param size the number of nodes + * @param numThreads threads to use (1, or fewer than {@value #PARALLEL_MIN_NODES} nodes = serial) * @return the NeighborArray */ - private NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size) { + private static NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size, int numThreads) + throws IOException { NeighborArray[] neighbors = new NeighborArray[size]; - for (int i = 0; i < size; i++) { - RowView rv = adjacency.getRow(i); + if (numThreads <= 1 + || size < PARALLEL_MIN_NODES + || (adjacency instanceof CuVSDeviceMatrix + && !fitsParallelGraphCopyBudget(adjacency.size(), adjacency.columns()))) { + fillNeighborRange(adjacency, neighbors, 0, size); + return neighbors; + } + CuVSMatrix source = adjacency; + CuVSHostMatrix hostCopy = null; + if (adjacency instanceof CuVSDeviceMatrix deviceAdjacency) { + hostCopy = deviceAdjacency.toHost(); + source = hostCopy; + } + try { + fillNeighborArrayParallel(source, neighbors, size, numThreads); + return neighbors; + } finally { + if (hostCopy != null) { + hostCopy.close(); + } + } + } + + /** Returns whether an INT32 adjacency can be copied without exceeding the native-host budget. */ + static boolean fitsParallelGraphCopyBudget(long rows, long columns) { + if (rows < 0 || columns < 0) { + return false; + } + if (rows == 0 || columns == 0) { + return true; + } + return rows <= MAX_PARALLEL_GRAPH_COPY_BYTES / Integer.BYTES / columns; + } + + /** + * Materializes disjoint node ranges concurrently. Each thread writes its own slots of {@code + * neighbors} and its own {@link NeighborArray} instances, so no synchronization is needed; {@code + * source} must be a host matrix (stateless {@code getRow}). + */ + private static void fillNeighborArrayParallel( + CuVSMatrix source, NeighborArray[] neighbors, int size, int numThreads) throws IOException { + ExecutorService pool = Executors.newFixedThreadPool(Math.max(1, numThreads - 1)); + try { + int perThread = (size + numThreads - 1) / numThreads; + List> tasks = new ArrayList<>(numThreads); + for (int t = 0; t < numThreads; t++) { + final int start = t * perThread; + final int end = Math.min(start + perThread, size); + if (start >= end) { + break; + } + tasks.add( + () -> { + fillNeighborRange(source, neighbors, start, end); + return null; + }); + } + new TaskExecutor(pool).invokeAll(tasks); + } finally { + pool.shutdown(); + } + } + + /** Fills {@code neighbors[start, end)} from the adjacency rows. */ + private static void fillNeighborRange( + CuVSMatrix source, NeighborArray[] neighbors, int start, int end) { + for (int i = start; i < end; i++) { + RowView rv = source.getRow(i); if (rv != null && rv.size() > 0) { - neighbors[i] = new NeighborArray((int) rv.size(), true); + NeighborArray na = new NeighborArray((int) rv.size(), true); for (int j = 0; j < rv.size(); j++) { - neighbors[i].addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); + na.addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); } + neighbors[i] = na; } else { neighbors[i] = new NeighborArray(0, true); } } - return neighbors; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index e3e10f44a1..5c22af2849 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -202,7 +202,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws dataset, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.NONE); + QuantizationType.NONE, + acceleratedHNSWParams.getWriterThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 10380ca538..1da1f9be77 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -183,7 +183,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw dataset, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.BINARY); + QuantizationType.BINARY, + acceleratedHNSWParams.getWriterThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index f9dbbfb06d..7e33bd79c4 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -208,7 +208,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE dataset, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.SCALAR); + QuantizationType.SCALAR, + acceleratedHNSWParams.getWriterThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java new file mode 100644 index 0000000000..492c437755 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java @@ -0,0 +1,232 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.RowView; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; +import org.junit.Test; + +/** Verifies serial and parallel CAGRA-adjacency materialization are equivalent and bounded. */ +public class TestWriterThreadsGraphMaterialization extends LuceneTestCase { + + private static final int NUM_NODES = GPUBuiltHnswGraph.PARALLEL_MIN_NODES + 1000; + private static final int DEGREE = 12; + private static final int NUM_THREADS = 4; + + @Test + public void parallelMaterializationMatchesSerial() throws Exception { + int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(1)); + + try (CuVSMatrix matrix = new ArrayMatrix(adjacency)) { + GPUBuiltHnswGraph serial = newSingleLayerGraph(matrix, 1); + GPUBuiltHnswGraph parallel = newSingleLayerGraph(matrix, NUM_THREADS); + assertGraphsEqual(serial, parallel); + } + } + + @Test + public void graphCopyBudgetHandlesExpectedDatasetSizesAndOverflow() { + assertTrue(GPUBuiltHnswGraph.fitsParallelGraphCopyBudget(25_000_000L, 32)); + assertFalse(GPUBuiltHnswGraph.fitsParallelGraphCopyBudget(100_000_000L, 32)); + assertFalse(GPUBuiltHnswGraph.fitsParallelGraphCopyBudget(Long.MAX_VALUE, Long.MAX_VALUE)); + } + + @Test + public void oversizedDeviceAdjacencyUsesSerialFallback() throws Exception { + int[][] adjacency = randomAdjacency(NUM_NODES, 1, new Random(0)); + long oversizedColumns = + GPUBuiltHnswGraph.MAX_PARALLEL_GRAPH_COPY_BYTES / Integer.BYTES / NUM_NODES + 1; + try (CuVSMatrix matrix = new ArrayDeviceMatrix(adjacency, oversizedColumns)) { + GPUBuiltHnswGraph graph = newSingleLayerGraph(matrix, NUM_THREADS); + assertEquals(NUM_NODES, graph.size()); + } + } + + private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, int numThreads) + throws IOException { + return new GPUBuiltHnswGraph( + NUM_NODES, + /* dimensions= */ 4, + Arrays.asList((int[]) null), + List.of(layer0Adjacency), + numThreads); + } + + private static void assertGraphsEqual(HnswGraph a, HnswGraph b) throws Exception { + assertEquals(a.numLevels(), b.numLevels()); + for (int level = 0; level < a.numLevels(); level++) { + int[] nodes = NodesIterator.getSortedNodes(a.getNodesOnLevel(level)); + for (int node : nodes) { + assertArrayEquals( + "node " + node + " at level " + level + " has different neighbors", + arcsOf(a, level, node), + arcsOf(b, level, node)); + } + } + } + + private static int[] arcsOf(HnswGraph graph, int level, int node) throws Exception { + graph.seek(level, node); + List arcs = new ArrayList<>(); + for (int n = graph.nextNeighbor(); n != NO_MORE_DOCS; n = graph.nextNeighbor()) { + arcs.add(n); + } + return arcs.stream().mapToInt(Integer::intValue).toArray(); + } + + private static int[][] randomAdjacency(int numNodes, int degree, Random random) { + int[][] adjacency = new int[numNodes][degree]; + for (int[] row : adjacency) { + for (int j = 0; j < degree; j++) { + row[j] = random.nextInt(numNodes); + } + } + return adjacency; + } + + private static class ArrayMatrix implements CuVSMatrix { + private final int[][] rows; + + ArrayMatrix(int[][] rows) { + this.rows = rows; + } + + @Override + public long size() { + return rows.length; + } + + @Override + public long columns() { + return rows.length == 0 ? 0 : rows[0].length; + } + + @Override + public DataType dataType() { + return DataType.INT; + } + + @Override + public RowView getRow(long row) { + return new ArrayRow(rows[Math.toIntExact(row)]); + } + + @Override + public void toArray(int[][] target) { + for (int i = 0; i < rows.length; i++) { + System.arraycopy(rows[i], 0, target[i], 0, rows[i].length); + } + } + + @Override + public void toArray(float[][] target) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(byte[][] target) { + throw new UnsupportedOperationException(); + } + + @Override + public void toHost(CuVSHostMatrix target) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSHostMatrix toHost() { + throw new UnsupportedOperationException(); + } + + @Override + public void toDevice(CuVSDeviceMatrix target, CuVSResources resources) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSDeviceMatrix toDevice(CuVSResources resources) { + throw new UnsupportedOperationException(); + } + + @Override + public void close() {} + } + + /** Reports an oversized device shape and fails if the guarded host-copy path is reached. */ + private static final class ArrayDeviceMatrix extends ArrayMatrix implements CuVSDeviceMatrix { + private final long reportedColumns; + + ArrayDeviceMatrix(int[][] rows, long reportedColumns) { + super(rows); + this.reportedColumns = reportedColumns; + } + + @Override + public long columns() { + return reportedColumns; + } + + @Override + public CuVSHostMatrix toHost() { + throw new AssertionError("oversized device adjacency must not be copied to host"); + } + } + + private static final class ArrayRow implements RowView { + private final int[] values; + + ArrayRow(int[] values) { + this.values = values; + } + + @Override + public long size() { + return values.length; + } + + @Override + public int getAsInt(long index) { + return values[Math.toIntExact(index)]; + } + + @Override + public float getAsFloat(long index) { + throw new UnsupportedOperationException(); + } + + @Override + public byte getAsByte(long index) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(int[] target) { + System.arraycopy(values, 0, target, 0, values.length); + } + + @Override + public void toArray(float[] target) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(byte[] target) { + throw new UnsupportedOperationException(); + } + } +} From 9cce529738cb993be83babc0ae6998844e5aeea7 Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Wed, 16 Sep 2026 08:43:11 +0000 Subject: [PATCH 03/21] Parallelize bounded HNSW graph serialization --- .../cuvs/lucene/AcceleratedHNSWUtils.java | 197 +++++++++++---- .../Lucene99AcceleratedHNSWVectorsWriter.java | 6 +- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 6 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 6 +- .../TestWriterThreadsGraphSerialization.java | 224 ++++++++++++++++++ 5 files changed, 393 insertions(+), 46 deletions(-) create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index d5dfebb486..46ed0a0d8e 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -19,8 +19,14 @@ import java.util.Random; import java.util.SortedSet; import java.util.TreeSet; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.TaskExecutor; +import org.apache.lucene.store.ByteBuffersDataOutput; +import org.apache.lucene.store.DataOutput; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.hnsw.HnswGraph; @@ -222,56 +228,167 @@ private static CuVSMatrix buildCagraGraphForSubset( * @return a 2D array of offsets * @throws IOException I/O Exceptions */ - public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) + public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int numThreads) throws IOException { - // write vectors' neighbors on each level into the vectorIndex file int countOnLevel0 = graph.size(); - int[][] offsets = new int[graph.numLevels()][]; - int[] scratch = new int[graph.maxConn() * 2]; - for (int level = 0; level < graph.numLevels(); level++) { + int numLevels = graph.numLevels(); + int[][] offsets = new int[numLevels][]; + + // Level 0 holds all nodes and dominates serialization cost. Each node's delta/VInt block is + // independent, so encode level 0 in parallel and concatenate the per-thread buffers serially in + // node order, in memory-bounded waves. Higher levels are tiny and stay serial. The on-disk + // bytes + // are identical to the fully-serial path (blocks in node order, offsets = per-node byte + // lengths). + // graph.maxConn() scans every layer-0 adjacency row (O(graph size)); compute it once here + // rather than per level/per task below. + int maxConn = graph.maxConn(); + + int[] level0Nodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(0)); + offsets[0] = new int[level0Nodes.length]; + if (numThreads > 1 && level0Nodes.length >= PARALLEL_MIN_NODES) { + writeLevel0Parallel( + graph, vectorIndex, level0Nodes, offsets[0], countOnLevel0, maxConn, numThreads); + } else { + writeLevelSerial(graph, vectorIndex, 0, level0Nodes, offsets[0], countOnLevel0, maxConn); + } + + for (int level = 1; level < numLevels; level++) { int[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); offsets[level] = new int[sortedNodes.length]; - int nodeOffsetId = 0; - - for (int node : sortedNodes) { - // Get node neighbors - NeighborArray neighbors = graph.getNeighbors(level, node); - // Get the size of the neighbor array - int size = neighbors.size(); - // Write size in VInt as the neighbors list is typically small - long offsetStart = vectorIndex.getFilePointer(); - // Get neighbors - int[] nnodes = neighbors.nodes(); - // Sort them - Arrays.sort(nnodes, 0, size); - // Now that we have sorted, do delta encoding to minimize the required bits to store the - // information - int actualSize = 0; - if (size > 0) { - scratch[0] = nnodes[0]; - actualSize = 1; - } - // De-duplication - for (int i = 1; i < size; i++) { - assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; - // Sorting step helps here - if (nnodes[i - 1] == nnodes[i]) { + writeLevelSerial( + graph, vectorIndex, level, sortedNodes, offsets[level], countOnLevel0, maxConn); + } + return offsets; + } + + /** Node count below which parallel level-0 serialization is not worth the overhead. */ + static final int PARALLEL_MIN_NODES = 1 << 16; + + /** Maximum encoded payload held by one parallel wave before it is copied to the index output. */ + static final long MAX_PARALLEL_ENCODE_BYTES = 64L << 20; + + /** Maximum bytes written by one node: one VInt count and {@code maxConn} VInt deltas. */ + private static final int MAX_VINT_BYTES = 5; + + /** Serially encodes a level's nodes into {@code out}, recording per-node byte lengths. */ + private static void writeLevelSerial( + GPUBuiltHnswGraph graph, + IndexOutput out, + int level, + int[] sortedNodes, + int[] offsets, + int countOnLevel0, + int maxConn) + throws IOException { + int[] scratch = new int[maxConn * 2]; + int idx = 0; + for (int node : sortedNodes) { + long start = out.getFilePointer(); + encodeNode(graph.getNeighbors(level, node), scratch, out, countOnLevel0); + offsets[idx++] = Math.toIntExact(out.getFilePointer() - start); + } + } + + /** + * Encodes level 0 in parallel: within memory-bounded waves, threads encode contiguous node + * sub-ranges into per-thread buffers, which are then concatenated to {@code out} in node order + * (identical layout to the serial path). + */ + private static void writeLevel0Parallel( + GPUBuiltHnswGraph graph, + IndexOutput out, + int[] nodes, + int[] offsets, + int countOnLevel0, + int maxConn, + int numThreads) + throws IOException { + // invokeAll joins every task before it returns, including tasks still running when a sibling + // fails, so `buffers` is never concatenated while a worker might still be writing into it. + // It also runs one share of each wave on the calling thread instead of parking it, so the + // pool only has to cover the other ranges. + ExecutorService pool = Executors.newFixedThreadPool(Math.max(1, numThreads - 1)); + try { + TaskExecutor executor = new TaskExecutor(pool); + int n = nodes.length; + int waveNodes = nodesPerSerializationWave(maxConn); + for (int waveStart = 0; waveStart < n; ) { + int waveEnd = (int) Math.min(n, (long) waveStart + waveNodes); + int perThread = (waveEnd - waveStart + numThreads - 1) / numThreads; + + ByteBuffersDataOutput[] buffers = new ByteBuffersDataOutput[numThreads]; + List> tasks = new ArrayList<>(numThreads); + for (int t = 0; t < numThreads; t++) { + final int subStart = waveStart + t * perThread; + final int subEnd = Math.min(subStart + perThread, waveEnd); + final int slot = t; + if (subStart >= subEnd) { continue; } - scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; + tasks.add( + () -> { + ByteBuffersDataOutput buffer = new ByteBuffersDataOutput(); + int[] scratch = new int[maxConn * 2]; + for (int i = subStart; i < subEnd; i++) { + long before = buffer.size(); + encodeNode(graph.getNeighbors(0, nodes[i]), scratch, buffer, countOnLevel0); + offsets[i] = Math.toIntExact(buffer.size() - before); + } + buffers[slot] = buffer; + return null; + }); } - // Write the size after duplicates are removed - vectorIndex.writeVInt(actualSize); - // Write de-duplicated neighbors - for (int i = 0; i < actualSize; i++) { - vectorIndex.writeVInt(scratch[i]); + executor.invokeAll(tasks); + // Concatenate in thread order (== node order), preserving the serial byte layout. + for (ByteBuffersDataOutput buffer : buffers) { + if (buffer != null) { + buffer.copyTo(out); + } } - offsets[level][nodeOffsetId++] = - Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); + waveStart = waveEnd; } + } finally { + pool.shutdown(); + } + } + + /** + * Sizes serialization waves from an upper bound on encoded bytes rather than only node count, so + * increasing graph degree cannot create an unbounded transient heap allocation. + */ + static int nodesPerSerializationWave(int maxConn) { + long maxBytesPerNode = (Math.max(0L, maxConn) + 1L) * MAX_VINT_BYTES; + return (int) + Math.max(1L, Math.min(Integer.MAX_VALUE, MAX_PARALLEL_ENCODE_BYTES / maxBytesPerNode)); + } + + /** + * Sorts, delta-encodes and de-duplicates a node's neighbors and writes the block (VInt size + VInt + * deltas) to {@code out}. Shared by the serial and parallel paths so encoding is identical. + */ + private static void encodeNode( + NeighborArray neighbors, int[] scratch, DataOutput out, int countOnLevel0) + throws IOException { + int size = neighbors == null ? 0 : neighbors.size(); + int actualSize = 0; + if (size > 0) { + int[] nnodes = neighbors.nodes(); + Arrays.sort(nnodes, 0, size); + scratch[0] = nnodes[0]; + actualSize = 1; + for (int i = 1; i < size; i++) { + assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; + if (nnodes[i - 1] == nnodes[i]) { + continue; + } + scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; + } + } + out.writeVInt(actualSize); + for (int i = 0; i < actualSize; i++) { + out.writeVInt(scratch[i]); } - // Return offsets (information written while writing the meta info) - return offsets; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 5c22af2849..1cf4fb50bd 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -205,7 +205,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws QuantizationType.NONE, acceleratedHNSWParams.getWriterThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -280,7 +281,8 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) int dimensions = fieldInfo.getVectorDimension(); GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 1da1f9be77..cf6d19f7ca 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -188,7 +188,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata @@ -276,7 +277,8 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 7e33bd79c4..aebfbea7c6 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -214,7 +214,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; @@ -301,7 +302,8 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java new file mode 100644 index 0000000000..5ddb59cf88 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java @@ -0,0 +1,224 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.RowView; +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.Test; + +/** Verifies parallel level-zero graph serialization is byte-identical and memory-bounded. */ +public class TestWriterThreadsGraphSerialization extends LuceneTestCase { + + private static final int NUM_NODES = AcceleratedHNSWUtils.PARALLEL_MIN_NODES + 1000; + private static final int DEGREE = 12; + private static final int REPORTED_MAX_CONN = 512; + private static final int NUM_THREADS = 4; + + @Test + public void parallelSerializationMatchesSerialAcrossWaves() throws Exception { + int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(2)); + + try (CuVSMatrix matrix = new ArrayMatrix(adjacency); + Directory dir = new ByteBuffersDirectory()) { + GPUBuiltHnswGraph serialGraph = newSingleLayerGraph(matrix); + GPUBuiltHnswGraph parallelGraph = newSingleLayerGraph(matrix); + + int[][] serialOffsets; + try (IndexOutput out = dir.createOutput("serial", IOContext.DEFAULT)) { + serialOffsets = AcceleratedHNSWUtils.writeGraph(serialGraph, out, 1); + } + int[][] parallelOffsets; + try (IndexOutput out = dir.createOutput("parallel", IOContext.DEFAULT)) { + parallelOffsets = AcceleratedHNSWUtils.writeGraph(parallelGraph, out, NUM_THREADS); + } + + assertEquals(serialOffsets.length, parallelOffsets.length); + for (int level = 0; level < serialOffsets.length; level++) { + assertArrayEquals(serialOffsets[level], parallelOffsets[level]); + } + assertArrayEquals(readAllBytes(dir, "serial"), readAllBytes(dir, "parallel")); + + assertTrue( + "test must cross a serialization-wave boundary", + AcceleratedHNSWUtils.nodesPerSerializationWave(REPORTED_MAX_CONN) < NUM_NODES); + } + } + + @Test + public void serializationWaveHonorsEncodedByteBudget() { + for (int degree : new int[] {1, 32, 88, 152, 512}) { + int nodes = AcceleratedHNSWUtils.nodesPerSerializationWave(degree); + long maximumEncodedBytes = (long) nodes * (degree + 1L) * 5L; + assertTrue(maximumEncodedBytes <= AcceleratedHNSWUtils.MAX_PARALLEL_ENCODE_BYTES); + assertTrue(nodes > 0); + } + } + + private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency) + throws IOException { + return new ReportedMaxConnGraph( + NUM_NODES, + /* dimensions= */ 4, + Arrays.asList((int[]) null), + List.of(layer0Adjacency), + REPORTED_MAX_CONN); + } + + private static byte[] readAllBytes(Directory dir, String name) throws Exception { + try (IndexInput in = dir.openInput(name, IOContext.DEFAULT)) { + byte[] bytes = new byte[(int) in.length()]; + in.readBytes(bytes, 0, bytes.length); + return bytes; + } + } + + private static int[][] randomAdjacency(int numNodes, int degree, Random random) { + int[][] adjacency = new int[numNodes][degree]; + for (int[] row : adjacency) { + for (int j = 0; j < degree; j++) { + row[j] = random.nextInt(numNodes); + } + } + return adjacency; + } + + /** Inflates maxConn only to force this modest test graph through multiple bounded waves. */ + private static final class ReportedMaxConnGraph extends GPUBuiltHnswGraph { + private final int reportedMaxConn; + + ReportedMaxConnGraph( + int size, + int dimensions, + List layerNodes, + List layerAdjacencies, + int reportedMaxConn) + throws IOException { + super(size, dimensions, layerNodes, layerAdjacencies, 1); + this.reportedMaxConn = reportedMaxConn; + } + + @Override + public int maxConn() { + return reportedMaxConn; + } + } + + private static final class ArrayMatrix implements CuVSMatrix { + private final int[][] rows; + + ArrayMatrix(int[][] rows) { + this.rows = rows; + } + + @Override + public long size() { + return rows.length; + } + + @Override + public long columns() { + return rows.length == 0 ? 0 : rows[0].length; + } + + @Override + public DataType dataType() { + return DataType.INT; + } + + @Override + public RowView getRow(long row) { + int[] values = rows[Math.toIntExact(row)]; + return new RowView() { + @Override + public long size() { + return values.length; + } + + @Override + public int getAsInt(long index) { + return values[Math.toIntExact(index)]; + } + + @Override + public float getAsFloat(long index) { + throw new UnsupportedOperationException(); + } + + @Override + public byte getAsByte(long index) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(int[] target) { + System.arraycopy(values, 0, target, 0, values.length); + } + + @Override + public void toArray(float[] target) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(byte[] target) { + throw new UnsupportedOperationException(); + } + }; + } + + @Override + public void toArray(int[][] target) { + for (int i = 0; i < rows.length; i++) { + System.arraycopy(rows[i], 0, target[i], 0, rows[i].length); + } + } + + @Override + public void toArray(float[][] target) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(byte[][] target) { + throw new UnsupportedOperationException(); + } + + @Override + public void toHost(CuVSHostMatrix target) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSHostMatrix toHost() { + throw new UnsupportedOperationException(); + } + + @Override + public void toDevice(CuVSDeviceMatrix target, CuVSResources resources) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSDeviceMatrix toDevice(CuVSResources resources) { + throw new UnsupportedOperationException(); + } + + @Override + public void close() {} + } +} From d3220f6cbfa0b13d1ad05056b1da8f568efda279 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 18 Sep 2026 10:39:55 +0000 Subject: [PATCH 04/21] Harden accelerated HNSW matrix ownership and merges --- .../com/nvidia/cuvs/CuVSDeviceMatrix.java | 20 +- .../java/com/nvidia/cuvs/CuVSHostMatrix.java | 20 +- .../main/java/com/nvidia/cuvs/CuVSMatrix.java | 19 +- .../com/nvidia/cuvs/internal/common/Util.java | 5 +- .../com/nvidia/cuvs/spi/JDKProvider.java | 38 +- .../java/com/nvidia/cuvs/CuVSMatrixIT.java | 34 + .../cuvs/lucene/AcceleratedHNSWUtils.java | 609 ++++++++++++++---- .../cuvs/lucene/CuVS2510GPUVectorsWriter.java | 171 ++++- .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 100 ++- .../Lucene99AcceleratedHNSWVectorsWriter.java | 82 ++- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 120 +++- ...ratedHNSWScalarQuantizedVectorsWriter.java | 143 ++-- .../java/com/nvidia/cuvs/lucene/Utils.java | 113 +++- .../TestAcceleratedHNSWDeletedDocuments.java | 60 ++ ...estAcceleratedHNSWMultiLayerRoundTrip.java | 171 +++++ .../lucene/TestAcceleratedHNSWQuantizers.java | 55 ++ .../TestAcceleratedHNSWUpperLayers.java | 273 ++++++++ ...VS2510GPUVectorsWriterFailureHandling.java | 150 +++++ .../lucene/TestQuantizedVectorsFormats.java | 121 ++++ .../lucene/TestUtilsThrowableHandling.java | 69 ++ 20 files changed, 2045 insertions(+), 328 deletions(-) create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWQuantizers.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUpperLayers.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsWriterFailureHandling.java diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java index 16b1881ad5..73fc0519a8 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java @@ -15,8 +15,22 @@ public interface CuVSDeviceMatrix extends CuVSMatrix { * responsible to call {@link CuVSMatrix#close()} to free its resources when done. */ default CuVSHostMatrix toHost() { - var hostMatrix = CuVSMatrix.hostBuilder(size(), columns(), dataType()).build(); - toHost(hostMatrix); - return hostMatrix; + CuVSHostMatrix hostMatrix; + try (var builder = CuVSMatrix.hostBuilder(size(), columns(), dataType())) { + hostMatrix = builder.build(); + } + try { + toHost(hostMatrix); + return hostMatrix; + } catch (RuntimeException | Error failure) { + try { + hostMatrix.close(); + } catch (RuntimeException | Error closeFailure) { + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } + } + throw failure; + } } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java index 62b708f71a..3607fe0130 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java @@ -11,8 +11,22 @@ public interface CuVSHostMatrix extends CuVSMatrix { int get(int row, int col); default CuVSDeviceMatrix toDevice(CuVSResources resources) { - var deviceMatrix = CuVSMatrix.deviceBuilder(resources, size(), columns(), dataType()).build(); - toDevice(deviceMatrix, resources); - return deviceMatrix; + CuVSDeviceMatrix deviceMatrix; + try (var builder = CuVSMatrix.deviceBuilder(resources, size(), columns(), dataType())) { + deviceMatrix = builder.build(); + } + try { + toDevice(deviceMatrix, resources); + return deviceMatrix; + } catch (RuntimeException | Error failure) { + try { + deviceMatrix.close(); + } catch (RuntimeException | Error closeFailure) { + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } + } + throw failure; + } } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java index 097efb1003..4447f88db6 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java @@ -73,7 +73,7 @@ static CuVSMatrix ofArray(byte[][] vectors) { * A builder to construct a new matrix one row at a time * @param the CuVSMatrix type to build */ - interface Builder { + interface Builder extends AutoCloseable { /** * Adds a single vector to the matrix. * @@ -102,7 +102,24 @@ interface Builder { */ void addVector(short[] vector); + /** + * Completes the matrix and transfers ownership to the caller. + * + *

If this method fails, closing the builder releases any matrix storage allocated while the + * builder was created. + */ T build(); + + /** + * Closes this builder. Built-in builders release matrix storage unless ownership was + * transferred by a successful {@link #build()}. + * + *

The default implementation preserves compatibility with providers compiled before + * builders became closeable. Builders that allocate storage before {@link #build()} should + * override this method. + */ + @Override + default void close() {} } /** diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java index 986ca5bd88..8025e9c640 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java @@ -11,6 +11,7 @@ import static com.nvidia.cuvs.internal.panama.headers_h.*; import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.LibraryException; import com.nvidia.cuvs.internal.panama.DLDataType; import com.nvidia.cuvs.internal.panama.DLDevice; import com.nvidia.cuvs.internal.panama.DLManagedTensor; @@ -126,7 +127,7 @@ public static int cudaGetDeviceProperties(MemorySegment prop, int device) { public static void checkCuVSError(int value, String caller) { if (value != CUVS_SUCCESS) { String errorMsg = getLastErrorText(); - throw new RuntimeException(caller + " returned " + value + "[" + errorMsg + "]"); + throw new LibraryException(caller + " returned " + value + "[" + errorMsg + "]"); } } @@ -138,7 +139,7 @@ public static void checkCuVSError(int value, String caller) { */ public static void checkCudaError(int value, String caller) { if (value != CUDA_SUCCESS) { - throw new RuntimeException(caller + " returned " + value); + throw new LibraryException(caller + " returned " + value); } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 1f16a4e904..c240e35d3e 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -617,7 +617,8 @@ public CuVSMatrix newMatrixFromArray(byte[][] vectors) { return dataset; } - private abstract static class MatrixBuilder { + private abstract static class MatrixBuilder + implements AutoCloseable { protected final long columns; protected final long size; @@ -626,6 +627,7 @@ private abstract static class MatrixBuilder { protected final long rowSize; protected final long rowBytes; protected int currentRow; + private boolean closed; protected MatrixBuilder(T matrix, long size, long columns) { this.columns = columns; @@ -635,6 +637,7 @@ protected MatrixBuilder(T matrix, long size, long columns) { this.rowSize = columns * elementSize; this.rowBytes = rowSize; this.currentRow = 0; + this.closed = false; } protected MatrixBuilder(T matrix, long size, long columns, int rowStride) { @@ -646,9 +649,11 @@ protected MatrixBuilder(T matrix, long size, long columns, int rowStride) { this.rowBytes = columns * elementSize; this.currentRow = 0; + this.closed = false; } public void addVector(float[] vector) { + ensureOpen(); if (vector.length != columns) { throw new IllegalArgumentException( String.format( @@ -658,6 +663,7 @@ public void addVector(float[] vector) { } public void addVector(byte[] vector) { + ensureOpen(); if (vector.length != columns) { throw new IllegalArgumentException( String.format( @@ -667,6 +673,7 @@ public void addVector(byte[] vector) { } public void addVector(int[] vector) { + ensureOpen(); if (vector.length != columns) { throw new IllegalArgumentException( String.format( @@ -676,6 +683,7 @@ public void addVector(int[] vector) { } public void addVector(short[] vector) { + ensureOpen(); if (vector.length != columns) { throw new IllegalArgumentException( String.format( @@ -684,6 +692,27 @@ public void addVector(short[] vector) { internalAddVector(MemorySegment.ofArray(vector)); } + protected final T transferOwnership() { + ensureOpen(); + closed = true; + return matrix; + } + + protected final void ensureOpen() { + if (closed) { + throw new IllegalStateException("matrix builder is closed"); + } + } + + @Override + public final void close() { + if (closed) { + return; + } + closed = true; + matrix.close(); + } + protected abstract void internalAddVector(MemorySegment vector); } @@ -772,11 +801,12 @@ private void flushBuffer(MemorySegment hostBuffer) { @Override public CuVSDeviceMatrix build() { + ensureOpen(); try (var access = resources.access()) { var hostBuffer = CuVSResourcesImpl.getHostBuffer(access); flushBuffer(hostBuffer); } - return matrix; + return transferOwnership(); } } @@ -831,7 +861,7 @@ protected void internalAddVector(MemorySegment vector) { @Override public CuVSDeviceMatrix build() { - return matrix; + return transferOwnership(); } } @@ -861,7 +891,7 @@ protected void internalAddVector(MemorySegment vector) { @Override public CuVSHostMatrix build() { - return matrix; + return transferOwnership(); } } } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java index 954a7d1e4f..4c9f8059b2 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java @@ -7,6 +7,7 @@ import static com.carrotsearch.randomizedtesting.RandomizedTest.*; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; import com.carrotsearch.randomizedtesting.RandomizedRunner; import com.nvidia.cuvs.spi.CuVSProvider; @@ -309,6 +310,39 @@ public void testFloatDatasetDeviceBuilder() throws Throwable { } } + @Test + public void testClosingDeviceBuilderReleasesUnbuiltMatrix() throws Throwable { + try (var resources = CheckedCuVSResources.create()) { + var builder = CuVSMatrix.deviceBuilder(resources, 2, 4, CuVSMatrix.DataType.FLOAT); + builder.addVector(new float[4]); + builder.close(); + builder.close(); + + assertThrows(IllegalStateException.class, builder::build); + assertThrows(IllegalStateException.class, () -> builder.addVector(new float[4])); + } + } + + @Test + public void testClosingDeviceBuilderAfterBuildDoesNotCloseMatrix() throws Throwable { + try (var resources = CheckedCuVSResources.create()) { + CuVSDeviceMatrix matrix; + try (var builder = CuVSMatrix.deviceBuilder(resources, 1, 4, CuVSMatrix.DataType.FLOAT)) { + builder.addVector(new float[] {1f, 2f, 3f, 4f}); + matrix = builder.build(); + + assertThrows(IllegalStateException.class, builder::build); + assertThrows(IllegalStateException.class, () -> builder.addVector(new float[4])); + } + + try (matrix) { + float[][] actual = new float[1][4]; + matrix.toArray(actual); + assertArrayEquals(new float[] {1f, 2f, 3f, 4f}, actual[0], DELTA); + } + } + } + private void testIntDatasetBuilder(int rows, int cols, CuVSMatrix.Builder builder) { var data = new int[rows][cols]; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 1be9f2450c..31c2213492 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -6,7 +6,6 @@ package com.nvidia.cuvs.lucene; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; -import static com.nvidia.cuvs.lucene.Utils.createByteMatrixFromArray; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraIndexParams; @@ -16,6 +15,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import java.util.Objects; import java.util.Random; import java.util.SortedSet; import java.util.TreeSet; @@ -58,35 +58,53 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens int[][] singleNodeAdjacency = new int[][] {{-1}}; // -1 indicates no neighbors // Create CuVSMatrix from the adjacency list - CuVSMatrix adjacencyMatrix = CuVSMatrix.ofArray(singleNodeAdjacency); - - // Create layer data for single-level graph - List layerNodes = new ArrayList<>(); - List layerAdjacencies = new ArrayList<>(); - - // Layer 0: contains all nodes (just the single node) - layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list - layerAdjacencies.add(adjacencyMatrix); - - // Create the single-layer graph - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + try (CuVSMatrix adjacencyMatrix = CuVSMatrix.ofArray(singleNodeAdjacency)) { + // Layer 0 contains every node, so its node list is implicit. + List layerNodes = new ArrayList<>(); + layerNodes.add(null); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, List.of(adjacencyMatrix)); + } } /** - * Creates a multi-layer HNSW graph with dynamic number of layers. - * M = ceil(cagraGraphDegree / 2), where cagraGraphDegree is the CAGRA adjacency list's degree - * (its column count). Ceil is used to accommodate odd graph degrees. - * Each layer contains 1/M nodes from the previous layer - * Creates layers until the highest layer has <= M nodes - *

- * Vectors for higher-layer subsets are read directly from the native matrix - * via {@link CuVSMatrix#getRow(long)} and {@link RowView#toArray(float[])}, - * avoiding any additional heap allocation of the full dataset. Used by both - * the flush and merge paths; the caller provides the vectors as a - * {@link CuVSMatrix}. + * Creates a multi-layer HNSW graph from heap vectors. + * + *

This overload preserves the original public API. Only rows selected for an upper layer are + * copied into native memory. */ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, + int size, + int dimensions, + CuVSMatrix adjacencyListMatrix, + List vectors, + int hnswLayers, + CagraIndexParams params, + QuantizationType quantization) + throws Throwable { + Objects.requireNonNull(vectors, "vectors"); + if (vectors.size() != size) { + throw new IllegalArgumentException( + "Expected " + size + " vectors, but received " + vectors.size()); + } + CuVSMatrix.DataType dataType = expectedDataType(quantization); + int columns = expectedColumns(dimensions, quantization); + return createMultiLayerHnswGraph( + size, + dimensions, + adjacencyListMatrix, + hnswLayers, + params, + selectedNodes -> createSubsetDataset(vectors, selectedNodes, columns, dataType)); + } + + /** + * Creates a multi-layer HNSW graph from a native matrix. + * + *

Only sampled rows are copied into each upper-layer matrix; the complete dataset is never + * materialized on the Java heap. + */ + static GPUBuiltHnswGraph createMultiLayerHnswGraph( int dimensions, CuVSMatrix adjacencyListMatrix, CuVSMatrix vectorDataset, @@ -94,9 +112,55 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( CagraIndexParams params, QuantizationType quantization) throws Throwable { + Objects.requireNonNull(vectorDataset, "vectorDataset"); + int size = Math.toIntExact(vectorDataset.size()); + CuVSMatrix.DataType dataType = expectedDataType(quantization); + int columns = expectedColumns(dimensions, quantization); + if (vectorDataset.columns() != columns) { + throw new IllegalArgumentException( + "Expected " + columns + " matrix columns, but received " + vectorDataset.columns()); + } + if (vectorDataset.dataType() != dataType) { + throw new IllegalArgumentException( + "Expected " + dataType + " matrix data, but received " + vectorDataset.dataType()); + } + return createMultiLayerHnswGraph( + size, + dimensions, + adjacencyListMatrix, + hnswLayers, + params, + selectedNodes -> createSubsetDataset(vectorDataset, selectedNodes, columns, dataType)); + } - int size = (int) vectorDataset.size(); - int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); + private static GPUBuiltHnswGraph createMultiLayerHnswGraph( + int size, + int dimensions, + CuVSMatrix adjacencyListMatrix, + int hnswLayers, + CagraIndexParams params, + SubsetDatasetFactory subsetDatasetFactory) + throws Throwable { + Objects.requireNonNull(adjacencyListMatrix, "adjacencyListMatrix"); + Objects.requireNonNull(params, "params"); + if (size < 2) { + throw new IllegalArgumentException("A multi-layer graph requires at least two vectors"); + } + if (adjacencyListMatrix.dataType() != CuVSMatrix.DataType.INT + && adjacencyListMatrix.dataType() != CuVSMatrix.DataType.UINT) { + throw new IllegalArgumentException( + "Expected INT or UINT adjacency data, but received " + adjacencyListMatrix.dataType()); + } + if (adjacencyListMatrix.size() != size) { + throw new IllegalArgumentException( + "Expected " + size + " adjacency rows, but received " + adjacencyListMatrix.size()); + } + + int degree = Math.toIntExact(adjacencyListMatrix.columns()); + if (degree <= 0) { + throw new IllegalArgumentException("The layer-0 graph must have a positive degree"); + } + int M = Math.ceilDiv(degree, 2); List layerNodes = new ArrayList<>(); List layerAdjacencies = new ArrayList<>(); @@ -109,108 +173,260 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( int layerIndex = 1; Random random = new Random(); - while (layerIndex < hnswLayers && currentLayerSize > 1) { - int nextLayerSize = Math.max(2, currentLayerSize / M); - SortedSet selectedNodesSet = new TreeSet<>(); + Throwable failure = null; + try { + while (layerIndex < hnswLayers && currentLayerSize > 1) { + int nextLayerSize = Math.max(2, currentLayerSize / M); + SortedSet selectedNodesSet = new TreeSet<>(); - if (layerIndex == 1) { - while (selectedNodesSet.size() < nextLayerSize) { - selectedNodesSet.add(random.nextInt(size)); + if (layerIndex == 1) { + while (selectedNodesSet.size() < nextLayerSize) { + selectedNodesSet.add(random.nextInt(size)); + } + } else { + int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); + while (selectedNodesSet.size() < nextLayerSize) { + selectedNodesSet.add(prevLayerNodes[random.nextInt(prevLayerNodes.length)]); + } } - } else { - int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); - while (selectedNodesSet.size() < nextLayerSize) { - selectedNodesSet.add(prevLayerNodes[random.nextInt(prevLayerNodes.length)]); + + int[] selectedNodes = + selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray(); + CuVSMatrix subsetDataset = subsetDatasetFactory.create(selectedNodes); + CuVSMatrix upperAdjacency = buildCagraGraphForSubset(subsetDataset, selectedNodes, params); + try { + layerNodes.add(selectedNodes); + layerAdjacencies.add(upperAdjacency); + } catch (Throwable addFailure) { + try { + upperAdjacency.close(); + } catch (Throwable closeFailure) { + if (addFailure != closeFailure) { + addFailure.addSuppressed(closeFailure); + } + } + throw addFailure; } - } - int[] selectedNodes = - selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray(); - layerNodes.add(selectedNodes); + currentLayerSize = nextLayerSize; + layerIndex++; + random = new Random(new Random().nextLong()); + } - if (quantization == QuantizationType.NONE) { - // Read only the sampled rows from the native matrix — no full-dataset heap copy - float[][] selectedVectors = new float[nextLayerSize][dimensions]; - for (int i = 0; i < nextLayerSize; i++) { - vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + } catch (Throwable t) { + failure = t; + throw t; + } finally { + Throwable closeFailure = null; + for (int i = layerAdjacencies.size() - 1; i >= 1; i--) { + try { + layerAdjacencies.get(i).close(); + } catch (Throwable t) { + if (closeFailure == null) { + closeFailure = t; + } else if (closeFailure != t) { + closeFailure.addSuppressed(t); + } } - layerAdjacencies.add( - buildCagraGraphForSubset( - selectedVectors, selectedNodes, 0, params, dimensions, quantization)); - } else { - // Byte width comes from the matrix itself: binary packs 8 dims/byte, scalar is 1 byte/dim. - int bytesPerVector = (int) vectorDataset.columns(); - byte[][] selectedVectors = new byte[nextLayerSize][bytesPerVector]; - for (int i = 0; i < nextLayerSize; i++) { - vectorDataset.getRow(selectedNodes[i]).toArray(selectedVectors[i]); + } + if (closeFailure != null) { + if (failure != null && failure != closeFailure) { + failure.addSuppressed(closeFailure); + } else { + throw closeFailure; } - layerAdjacencies.add( - buildCagraGraphForSubset( - selectedVectors, selectedNodes, bytesPerVector, params, dimensions, quantization)); } - - currentLayerSize = nextLayerSize; - layerIndex++; - random = new Random(new Random().nextLong()); } - - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); } - /** - * Builds a CAGRA graph for a subset of binary quantized vectors - */ + /** Builds a CAGRA graph for a sampled subset and remaps its ordinals to layer-0 ordinals. */ private static CuVSMatrix buildCagraGraphForSubset( - Object vectors, + CuVSMatrix subsetDataset, int[] selectedNodes, CagraIndexParams params) throws Throwable { + return buildCagraGraphForSubset( + subsetDataset, + selectedNodes, + params, + (dataset, indexParams) -> + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(dataset) + .withIndexParams(indexParams) + .build()); + } + + static CuVSMatrix buildCagraGraphForSubset( + CuVSMatrix subsetDataset, int[] selectedNodes, - int bytesPerVector, CagraIndexParams params, - int dimensions, - QuantizationType quantization) + SubsetIndexFactory indexFactory) throws Throwable { - - CuVSMatrix subsetDataset; - - if (quantization == QuantizationType.BINARY) { - subsetDataset = createByteMatrixFromArray((byte[][]) vectors, bytesPerVector); - } else if (quantization == QuantizationType.SCALAR) { - subsetDataset = createByteMatrixFromArray((byte[][]) vectors, dimensions); - } else { - subsetDataset = CuVSMatrix.ofArray((float[][]) vectors); + CuVSMatrix remappedAdjacency = null; + try (Utils.OwnedIndex ownedIndex = Utils.ownDataset(subsetDataset)) { + CagraIndex subsetIndex = indexFactory.build(subsetDataset, params); + ownedIndex.transferTo(subsetIndex); + remappedAdjacency = remapSubsetGraph(subsetIndex.getGraph(), selectedNodes); + } catch (Throwable failure) { + if (remappedAdjacency != null) { + try { + remappedAdjacency.close(); + } catch (Throwable closeFailure) { + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } + } + } + throw failure; } + return remappedAdjacency; + } - // Build CAGRA index for the subset - CagraIndex subsetIndex = - CagraIndex.newBuilder(getCuVSResourcesInstance()) - .withDataset(subsetDataset) - .withIndexParams(params) - .build(); - - // Get adjacency list from subset CAGRA index - CuVSMatrix cagraGraph = subsetIndex.getGraph(); + static CuVSMatrix remapSubsetGraph(CuVSMatrix cagraGraph, int[] selectedNodes) + throws IOException { + if (cagraGraph.dataType() != CuVSMatrix.DataType.INT + && cagraGraph.dataType() != CuVSMatrix.DataType.UINT) { + throw new IOException( + "Expected INT or UINT subset adjacency data, but received " + cagraGraph.dataType()); + } + if (cagraGraph.size() != selectedNodes.length) { + throw new IOException( + "Expected " + + selectedNodes.length + + " subset adjacency rows, but received " + + cagraGraph.size()); + } + int degree = Math.toIntExact(cagraGraph.columns()); + if (degree <= 0) { + throw new IOException("The subset graph must have a positive degree"); + } + int[] remappedRow = new int[degree]; + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(selectedNodes.length, degree, CuVSMatrix.DataType.INT)) { + for (int rowOrdinal = 0; rowOrdinal < selectedNodes.length; rowOrdinal++) { + RowView row = cagraGraph.getRow(rowOrdinal); + remapSubsetAdjacencyRow(rowOrdinal, row, degree, selectedNodes, remappedRow); + builder.addVector(remappedRow); + } + return builder.build(); + } + } - long numNodes = cagraGraph.size(); - long degree = cagraGraph.columns(); + static void remapSubsetAdjacencyRow( + int rowOrdinal, RowView source, int degree, int[] selectedNodes, int[] destination) + throws IOException { + if (source == null || source.size() != degree) { + throw new IOException( + "Expected " + + degree + + " neighbors for subset row " + + rowOrdinal + + ", but received " + + (source == null ? "null" : source.size())); + } + for (int column = 0; column < degree; column++) { + int subsetOrdinal = source.getAsInt(column); + if (subsetOrdinal < 0) { + destination[column] = subsetOrdinal; + } else if (subsetOrdinal >= selectedNodes.length) { + throw new IOException( + "Subset adjacency row " + + rowOrdinal + + ", column " + + column + + " contains local ordinal " + + subsetOrdinal + + " outside [0, " + + selectedNodes.length + + ")"); + } else { + destination[column] = selectedNodes[subsetOrdinal]; + } + } + } - // Create a re-mapped adjacency list - int[][] remappedAdjacency = new int[(int) numNodes][(int) degree]; + private static CuVSMatrix createSubsetDataset( + CuVSMatrix vectors, int[] selectedNodes, int columns, CuVSMatrix.DataType dataType) + throws IOException { + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(selectedNodes.length, columns, dataType)) { + if (dataType == CuVSMatrix.DataType.FLOAT) { + float[] rowBuffer = new float[columns]; + for (int node : selectedNodes) { + RowView row = vectors.getRow(node); + validateRowWidth(row, columns, node); + row.toArray(rowBuffer); + builder.addVector(rowBuffer); + } + } else { + byte[] rowBuffer = new byte[columns]; + for (int node : selectedNodes) { + RowView row = vectors.getRow(node); + validateRowWidth(row, columns, node); + row.toArray(rowBuffer); + builder.addVector(rowBuffer); + } + } + return builder.build(); + } + } - for (int i = 0; i < numNodes; i++) { - RowView rv = cagraGraph.getRow(i); - for (int j = 0; j < degree && j < rv.size(); j++) { - int subsetIndex1 = rv.getAsInt(j); - // Map subset index to original node ID - if (subsetIndex1 >= 0 && subsetIndex1 < selectedNodes.length) { - remappedAdjacency[i][j] = selectedNodes[subsetIndex1]; + private static CuVSMatrix createSubsetDataset( + List vectors, int[] selectedNodes, int columns, CuVSMatrix.DataType dataType) { + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(selectedNodes.length, columns, dataType)) { + for (int node : selectedNodes) { + Object vector = vectors.get(node); + if (dataType == CuVSMatrix.DataType.FLOAT) { + if (!(vector instanceof float[] values) || values.length != columns) { + throw new IllegalArgumentException( + "Vector " + node + " must be a float[" + columns + "]"); + } + builder.addVector(values); } else { - // Invalid index, use self-reference - remappedAdjacency[i][j] = selectedNodes[i]; + if (!(vector instanceof byte[] values) || values.length != columns) { + throw new IllegalArgumentException( + "Vector " + node + " must be a byte[" + columns + "]"); + } + builder.addVector(values); } } + return builder.build(); + } + } + + private static void validateRowWidth(RowView row, int columns, int rowIndex) throws IOException { + if (row == null || row.size() != columns) { + throw new IOException( + "Expected " + + columns + + " values for vector row " + + rowIndex + + ", but received " + + (row == null ? "null" : row.size())); + } + } + + private static int expectedColumns(int dimensions, QuantizationType quantization) { + if (dimensions <= 0) { + throw new IllegalArgumentException("Vector dimensions must be positive"); } + return quantization == QuantizationType.BINARY ? Math.ceilDiv(dimensions, 8) : dimensions; + } + + private static CuVSMatrix.DataType expectedDataType(QuantizationType quantization) { + Objects.requireNonNull(quantization, "quantization"); + return quantization == QuantizationType.NONE + ? CuVSMatrix.DataType.FLOAT + : CuVSMatrix.DataType.BYTE; + } + + @FunctionalInterface + private interface SubsetDatasetFactory { + CuVSMatrix create(int[] selectedNodes) throws Throwable; + } - subsetIndex.close(); - return CuVSMatrix.ofArray(remappedAdjacency); + @FunctionalInterface + interface SubsetIndexFactory { + CagraIndex build(CuVSMatrix dataset, CagraIndexParams params) throws Throwable; } /** @@ -399,30 +615,16 @@ public static List quantizeFloatVectorsToBinary(List floatVecto } int dimensions = floatVectors.get(0).length; - int numVectors = floatVectors.size(); - int bytesPerVector = (dimensions + 7) / 8; - - float[] centroids = new float[dimensions]; + BinaryQuantizer quantizer = new BinaryQuantizer(dimensions); for (float[] vector : floatVectors) { - for (int d = 0; d < dimensions; d++) { - centroids[d] += vector[d]; - } - } - for (int d = 0; d < dimensions; d++) { - centroids[d] /= numVectors; + quantizer.add(vector); } + quantizer.finish(); - List quantizedVectors = new ArrayList<>(numVectors); + List quantizedVectors = new ArrayList<>(floatVectors.size()); for (float[] vector : floatVectors) { - byte[] quantized = new byte[bytesPerVector]; - for (int d = 0; d < dimensions; d++) { - boolean bit = vector[d] > centroids[d]; - int byteIndex = d / 8; - int bitIndex = d % 8; - if (bit) { - quantized[byteIndex] |= (1 << bitIndex); - } - } + byte[] quantized = new byte[Math.ceilDiv(dimensions, 8)]; + quantizer.quantize(vector, quantized); quantizedVectors.add(quantized); } @@ -441,36 +643,159 @@ public static List quantizeFloatVectorsToScalar(List floatVecto } int dimensions = floatVectors.get(0).length; - int numVectors = floatVectors.size(); - - float[] minPerDim = new float[dimensions]; - float[] maxPerDim = new float[dimensions]; - Arrays.fill(minPerDim, Float.MAX_VALUE); - Arrays.fill(maxPerDim, Float.MIN_VALUE); - + ScalarQuantizer quantizer = new ScalarQuantizer(dimensions); for (float[] vector : floatVectors) { - for (int d = 0; d < dimensions; d++) { - minPerDim[d] = Math.min(minPerDim[d], vector[d]); - maxPerDim[d] = Math.max(maxPerDim[d], vector[d]); - } + quantizer.add(vector); } + quantizer.finish(); - List quantizedVectors = new ArrayList<>(numVectors); + List quantizedVectors = new ArrayList<>(floatVectors.size()); for (float[] vector : floatVectors) { byte[] quantized = new byte[dimensions]; - for (int d = 0; d < dimensions; d++) { - float range = maxPerDim[d] - minPerDim[d]; + quantizer.quantize(vector, quantized); + quantizedVectors.add(quantized); + } + + return quantizedVectors; + } + + static final class BinaryQuantizer { + private final float[] centroids; + private int count; + private boolean finished; + + BinaryQuantizer(int dimensions) { + if (dimensions <= 0) { + throw new IllegalArgumentException("Vector dimensions must be positive"); + } + centroids = new float[dimensions]; + } + + void add(float[] vector) { + ensureCollecting(); + requireDimensions(vector, centroids.length); + for (int dimension = 0; dimension < vector.length; dimension++) { + centroids[dimension] += vector[dimension]; + } + count = Math.incrementExact(count); + } + + void finish() { + ensureCollecting(); + if (count == 0) { + throw new IllegalStateException("Cannot finish binary quantization without vectors"); + } + for (int dimension = 0; dimension < centroids.length; dimension++) { + centroids[dimension] /= count; + } + finished = true; + } + + void quantize(float[] vector, byte[] destination) { + ensureFinished(); + requireDimensions(vector, centroids.length); + if (destination.length != Math.ceilDiv(centroids.length, 8)) { + throw new IllegalArgumentException("Binary quantization buffer dimensions do not match"); + } + Arrays.fill(destination, (byte) 0); + for (int dimension = 0; dimension < vector.length; dimension++) { + if (vector[dimension] > centroids[dimension]) { + destination[dimension / 8] |= (byte) (1 << (dimension % 8)); + } + } + } + + int count() { + return count; + } + + private void ensureCollecting() { + if (finished) { + throw new IllegalStateException("Quantization statistics are already final"); + } + } + + private void ensureFinished() { + if (finished == false) { + throw new IllegalStateException("Quantization statistics are not final"); + } + } + } + + static final class ScalarQuantizer { + private final float[] minima; + private final float[] maxima; + private int count; + private boolean finished; + + ScalarQuantizer(int dimensions) { + if (dimensions <= 0) { + throw new IllegalArgumentException("Vector dimensions must be positive"); + } + minima = new float[dimensions]; + maxima = new float[dimensions]; + Arrays.fill(minima, Float.MAX_VALUE); + // Preserve the existing scalar quantizer's behavior for all-negative dimensions. + Arrays.fill(maxima, Float.MIN_VALUE); + } + + void add(float[] vector) { + ensureCollecting(); + requireDimensions(vector, minima.length); + for (int dimension = 0; dimension < vector.length; dimension++) { + minima[dimension] = Math.min(minima[dimension], vector[dimension]); + maxima[dimension] = Math.max(maxima[dimension], vector[dimension]); + } + count = Math.incrementExact(count); + } + + void finish() { + ensureCollecting(); + if (count == 0) { + throw new IllegalStateException("Cannot finish scalar quantization without vectors"); + } + finished = true; + } + + void quantize(float[] vector, byte[] destination) { + ensureFinished(); + requireDimensions(vector, minima.length); + if (destination.length != minima.length) { + throw new IllegalArgumentException("Scalar quantization buffer dimensions do not match"); + } + for (int dimension = 0; dimension < vector.length; dimension++) { + float range = maxima[dimension] - minima[dimension]; if (range > 0) { - float normalized = (vector[d] - minPerDim[d]) / range; + float normalized = (vector[dimension] - minima[dimension]) / range; int quantizedValue = Math.round(normalized * 127.0f) - 64; - quantized[d] = (byte) Math.max(-64, Math.min(63, quantizedValue)); + destination[dimension] = (byte) Math.max(-64, Math.min(63, quantizedValue)); } else { - quantized[d] = 0; + destination[dimension] = 0; } } - quantizedVectors.add(quantized); } - return quantizedVectors; + int count() { + return count; + } + + private void ensureCollecting() { + if (finished) { + throw new IllegalStateException("Quantization statistics are already final"); + } + } + + private void ensureFinished() { + if (finished == false) { + throw new IllegalStateException("Quantization statistics are not final"); + } + } + } + + private static void requireDimensions(float[] vector, int dimensions) { + if (vector.length != dimensions) { + throw new IllegalArgumentException( + "Expected " + dimensions + " vector dimensions, but received " + vector.length); + } } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 3a7eb7852e..4460ee9187 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -21,6 +21,7 @@ import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.LibraryException; import java.io.IOException; import java.io.OutputStream; import java.nio.file.Files; @@ -204,17 +205,15 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro CuVSMatrix cagraDataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); writeCagraIndex(cagraIndexOutputStream, cagraDataset); - } catch (Throwable t) { - // Fallback to brute force in a few cases, for now. - // Log it to make it more obvious that this is what is happening. + } catch (RecoverableCagraConstructionException recoverable) { + ensureFallbackOutputUnchanged(cagraIndexOffset, cuvsIndex.getFilePointer(), recoverable); info( infoStream, COMPONENT, "CAGRA build failed for field \"" + fieldInfo.name + "\", falling back to a brute force index: " - + t); - Utils.handleThrowableWithIgnore(t, t.getMessage()); + + recoverable.getCause()); indexType = IndexType.BRUTE_FORCE; } cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; @@ -248,30 +247,47 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro * @throws Throwable */ private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { - CagraIndexParams params = - CagraIndexParamsFactory.create(gpuSearchParams, dataset.size(), dataset.columns()); - try (CagraIndex index = - CagraIndex.newBuilder(getCuVSResourcesInstance()) - .withDataset(dataset) - .withIndexParams(params) - .build(); - var deviceVectors = dataset.toDevice(getCuVSResourcesInstance())) { + CagraIndex index = null; + CuVSMatrix deviceVectors = null; + AutoCloseable indexDataset = null; + Throwable failure = null; + CagraWriteContext writeContext = new CagraWriteContext(); + try { + CagraIndexParams params = + CagraIndexParamsFactory.create(gpuSearchParams, dataset.size(), dataset.columns()); + index = + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(dataset) + .withIndexParams(params) + .build(); + deviceVectors = dataset.toDevice(getCuVSResourcesInstance()); /* * cuVS rejects makePaddedDataset for a device matrix whose rows already sit at the required - * stride, and asks for a view over that storage instead. Copying would be pointless there - * anyway, so pick the factory that matches the layout. + * stride, and asks for a view over that storage instead. */ if (CagraIndex.isPaddedDataset(deviceVectors)) { - try (var indexDatasetView = index.makePaddedDatasetView(deviceVectors)) { - index.updateDataset(indexDatasetView); - index.serialize(os); - } + var indexDatasetView = index.makePaddedDatasetView(deviceVectors); + indexDataset = indexDatasetView; + index.updateDataset(indexDatasetView); } else { - try (var indexDataset = index.makePaddedDataset(deviceVectors)) { - index.updateDataset(indexDataset); - index.serialize(os); - } + var paddedDataset = index.makePaddedDataset(deviceVectors); + indexDataset = paddedDataset; + index.updateDataset(paddedDataset); } + + // Invoking the serializer is the persistence boundary. No failure from this point can fall + // back safely, even if a particular serializer implementation happens not to write bytes. + writeContext.beginPersistence(); + index.serialize(os); + } catch (Throwable t) { + failure = t; + } + + Throwable cleanupFailure = closeCagraResources(index, dataset, indexDataset, deviceVectors); + failure = combineOperationAndCleanupFailures(failure, cleanupFailure); + if (failure != null) { + throw classifyCagraWriteFailure( + failure, writeContext.persistenceStarted(), cleanupFailure != null); } } @@ -287,13 +303,108 @@ private void writeBruteForceIndex(OutputStream os, CuVSMatrix dataset) throws Th new BruteForceIndexParams.Builder() .withNumWriterThreads(gpuSearchParams.getWriterThreads()) .build(); - var index = - BruteForceIndex.newBuilder(getCuVSResourcesInstance()) - .withIndexParams(params) - .withDataset(dataset) - .build(); - index.serialize(os); - index.close(); + try (dataset; + var index = + BruteForceIndex.newBuilder(getCuVSResourcesInstance()) + .withIndexParams(params) + .withDataset(dataset) + .build()) { + index.serialize(os); + } + } + + private static Throwable closeResource(AutoCloseable resource, Throwable failure) { + if (resource == null) { + return failure; + } + try { + resource.close(); + } catch (Throwable closeFailure) { + return addFailure(failure, closeFailure); + } + return failure; + } + + /** Closes CAGRA resources in dependency order and returns the first cleanup failure. */ + static Throwable closeCagraResources( + AutoCloseable index, + AutoCloseable originalDataset, + AutoCloseable indexDataset, + AutoCloseable deviceVectors) { + Throwable failure = null; + try { + if (index == null) { + originalDataset.close(); + } else { + Utils.closeIndexWithDatasetFallback(index, originalDataset); + } + } catch (Throwable closeFailure) { + failure = addFailure(failure, closeFailure); + } + failure = closeResource(indexDataset, failure); + return closeResource(deviceVectors, failure); + } + + static Throwable combineOperationAndCleanupFailures( + Throwable operationFailure, Throwable cleanupFailure) { + return addFailure(operationFailure, cleanupFailure); + } + + private static Throwable addFailure(Throwable primary, Throwable secondary) { + if (secondary == null) { + return primary; + } + if (primary == null) { + return secondary; + } + if (primary != secondary) { + primary.addSuppressed(secondary); + } + return primary; + } + + /** Signals the only failure for which the caller may safely build a brute-force index instead. */ + static final class RecoverableCagraConstructionException extends Exception { + private static final long serialVersionUID = 1L; + + RecoverableCagraConstructionException(LibraryException cause) { + super(cause); + } + } + + static Throwable classifyCagraWriteFailure( + Throwable failure, boolean persistenceStarted, boolean cleanupFailed) { + if (failure instanceof LibraryException libraryFailure + && persistenceStarted == false + && cleanupFailed == false + && failure.getSuppressed().length == 0) { + return new RecoverableCagraConstructionException(libraryFailure); + } + return failure; + } + + static void ensureFallbackOutputUnchanged(long expected, long actual, Throwable failure) + throws IOException { + if (actual != expected) { + throw new IOException( + "Cannot fall back after the CAGRA output position changed from " + + expected + + " to " + + actual, + failure); + } + } + + static final class CagraWriteContext { + private boolean persistenceStarted; + + void beginPersistence() { + persistenceStarted = true; + } + + boolean persistenceStarted() { + return persistenceStarted; + } } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index 7e9f888e32..fff21b62fa 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -9,6 +9,8 @@ import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.RowView; import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; import java.util.List; import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.NeighborArray; @@ -27,6 +29,7 @@ public class GPUBuiltHnswGraph extends HnswGraph { // Store layers data - each layer has its own nodes and adjacency lists private final List layerNodes; private final List layerNeighbors; + private final boolean[] sortedUpperLayers; // Layer 0 is special - it contains all nodes private final NeighborArray[] layer0Neighbors; @@ -41,12 +44,22 @@ public class GPUBuiltHnswGraph extends HnswGraph { */ public GPUBuiltHnswGraph( int size, int dimensions, List layerNodes, List layerAdjacencies) { - + if (size < 0) { + throw new IllegalArgumentException("Graph size must not be negative"); + } + if (dimensions <= 0) { + throw new IllegalArgumentException("Vector dimensions must be positive"); + } + if (layerAdjacencies.isEmpty() || layerNodes.size() != layerAdjacencies.size()) { + throw new IllegalArgumentException( + "Layer node and adjacency lists must have the same non-zero size"); + } this.size = size; this.dimensions = dimensions; this.numLevels = layerAdjacencies.size(); this.layerNodes = new ArrayList<>(); this.layerNeighbors = new ArrayList<>(); + this.sortedUpperLayers = new boolean[Math.max(0, numLevels - 1)]; // Process Layer 0 (base layer with all nodes) CuVSMatrix layer0Adjacency = layerAdjacencies.get(0); @@ -54,7 +67,12 @@ public GPUBuiltHnswGraph( // Process higher layers (1 to numLevels-1) for (int level = 1; level < numLevels; level++) { - int[] nodes = layerNodes.get(level); + int[] suppliedNodes = layerNodes.get(level); + if (suppliedNodes == null) { + throw new IllegalArgumentException("Missing node ordinals for level " + level); + } + int[] nodes = suppliedNodes.clone(); + sortedUpperLayers[level - 1] = validateLayerNodes(nodes, level); CuVSMatrix adjacency = layerAdjacencies.get(level); this.layerNodes.add(nodes); this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length)); @@ -69,21 +87,75 @@ public GPUBuiltHnswGraph( * @return the NeighborArray */ private NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size) { + if (adjacency.dataType() != CuVSMatrix.DataType.INT + && adjacency.dataType() != CuVSMatrix.DataType.UINT) { + throw new IllegalArgumentException( + "Expected INT or UINT adjacency data, but received " + adjacency.dataType()); + } + if (adjacency.size() != size) { + throw new IllegalArgumentException( + "Expected " + size + " adjacency rows, but received " + adjacency.size()); + } + int degree = Math.toIntExact(adjacency.columns()); + if (degree <= 0) { + throw new IllegalArgumentException("Adjacency matrices must have a positive degree"); + } NeighborArray[] neighbors = new NeighborArray[size]; for (int i = 0; i < size; i++) { RowView rv = adjacency.getRow(i); - if (rv != null && rv.size() > 0) { - neighbors[i] = new NeighborArray((int) rv.size(), true); - for (int j = 0; j < rv.size(); j++) { - neighbors[i].addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); + if (rv == null || rv.size() != degree) { + throw new IllegalArgumentException( + "Expected " + + degree + + " neighbors for adjacency row " + + i + + ", but received " + + (rv == null ? "null" : rv.size())); + } + neighbors[i] = new NeighborArray(degree, true); + for (int j = 0; j < degree; j++) { + int neighbor = rv.getAsInt(j); + if (neighbor < 0) { + continue; + } + if (neighbor >= this.size) { + throw new IllegalArgumentException( + "Adjacency row " + + i + + " contains ordinal " + + neighbor + + " outside graph size " + + this.size); } - } else { - neighbors[i] = new NeighborArray(0, true); + neighbors[i].addInOrder(neighbor, 1.0f - (j * 0.001f)); } } return neighbors; } + private boolean validateLayerNodes(int[] nodes, int level) { + boolean sorted = true; + int previous = -1; + for (int node : nodes) { + if (node < 0 || node >= size) { + throw new IllegalArgumentException( + "Level " + level + " contains node ordinal " + node + " outside [0, " + size + ")"); + } + sorted &= node > previous; + previous = node; + } + if (sorted == false) { + var uniqueNodes = new HashSet(nodes.length); + for (int node : nodes) { + if (uniqueNodes.add(node) == false) { + throw new IllegalArgumentException( + "Level " + level + " contains duplicate node ordinal " + node); + } + } + } + return sorted; + } + /** * Get all nodes on a given level as node 0th ordinals. */ @@ -111,11 +183,13 @@ public NeighborArray getNeighbors(int level, int node) { } else if (level > 0 && level < numLevels) { int[] nodes = layerNodes.get(level - 1); NeighborArray[] neighbors = layerNeighbors.get(level - 1); - - // Find the index of this node in the layer - for (int i = 0; i < nodes.length; i++) { - if (nodes[i] == node) { - return neighbors[i]; + if (sortedUpperLayers[level - 1]) { + int index = Arrays.binarySearch(nodes, node); + return index >= 0 ? neighbors[index] : null; + } + for (int index = 0; index < nodes.length; index++) { + if (nodes[index] == node) { + return neighbors[index]; } } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index e3e10f44a1..bc7124c409 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -150,12 +150,7 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException * @throws IOException */ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { - if (vectors.size() == 0) { - writeEmpty(fieldInfo, hnswMeta); - return; - } - if (vectors.size() < 2) { - writeSingleVectorGraph(fieldInfo, vectors); + if (writeTrivialField(fieldInfo, vectors.size())) { return; } CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); @@ -173,18 +168,11 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro * @throws IOException */ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { - int size = (int) dataset.size(); - if (size == 0) { - writeEmpty(fieldInfo, hnswMeta); - return; - } - if (size < 2) { - float[] buf = new float[fieldInfo.getVectorDimension()]; - dataset.getRow(0).toArray(buf); - writeSingleVectorGraph(fieldInfo, List.of(buf)); - return; - } - try { + try (Utils.OwnedIndex ownedIndex = Utils.ownDataset(dataset)) { + int size = Math.toIntExact(dataset.size()); + if (writeTrivialField(fieldInfo, size)) { + return; + } CagraIndexParams params = CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); CagraIndex cagraIndex = @@ -192,11 +180,11 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws .withDataset(dataset) .withIndexParams(params) .build(); + ownedIndex.transferTo(cagraIndex); CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); int dimensions = fieldInfo.getVectorDimension(); GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( - fieldInfo, dimensions, adjacencyListMatrix, dataset, @@ -215,12 +203,24 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws size, hnswGraph, graphLevelNodeOffsets); - cagraIndex.close(); } catch (Throwable t) { - Utils.handleThrowable(t); + throw Utils.handleThrowable(t); } } + /** Writes the empty or one-vector representation, if {@code size} is trivial. */ + private boolean writeTrivialField(FieldInfo fieldInfo, int size) throws IOException { + if (size == 0) { + writeEmpty(fieldInfo, hnswMeta); + return true; + } + if (size == 1) { + writeSingleVectorGraph(fieldInfo); + return true; + } + return false; + } + /** * Build the indexes and writes it to the disk. */ @@ -269,11 +269,9 @@ private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) thr * Builds and writes a single vector graph. * * @param fieldInfo instance of FieldInfo - * @param vectors the list of float vectors * @throws IOException I/O Exceptions */ - private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) - throws IOException { + private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { try { int size = 1; int dimensions = fieldInfo.getVectorDimension(); @@ -303,21 +301,37 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) */ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { try { + int size = countMergedVectors(fieldInfo, mergeState); + if (writeTrivialField(fieldInfo, size)) { + return; + } + int dims = fieldInfo.getVectorDimension(); FloatVectorValues mergedVectors = KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); - int size = mergedVectors.size(); - int dims = fieldInfo.getVectorDimension(); - CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT); - KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); - for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { - builder.addVector(mergedVectors.vectorValue(it.index())); + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT)) { + KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); + for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { + builder.addVector(mergedVectors.vectorValue(it.index())); + } + writeFieldInternal(fieldInfo, builder.build()); } - CuVSHostMatrix dataset = builder.build(); - writeFieldInternal(fieldInfo, dataset); } catch (Throwable t) { - Utils.handleThrowable(t); + throw Utils.handleThrowable(t); + } + } + + /** Counts the live vectors that the merge iterator will actually yield. */ + private static int countMergedVectors(FieldInfo fieldInfo, MergeState mergeState) + throws IOException { + FloatVectorValues mergedVectors = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + int count = 0; + KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); + for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { + count = Math.incrementExact(count); } + return count; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 10380ca538..4bdf8a21cc 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -7,7 +7,6 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createMultiLayerHnswGraph; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createSingleVectorHnswGraph; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.printInfoStream; -import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToBinary; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeEmpty; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeGraph; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeMeta; @@ -24,6 +23,7 @@ import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.BinaryQuantizer; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; import java.io.IOException; import java.util.ArrayList; @@ -146,38 +146,45 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException * @throws IOException */ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { - if (vectors.size() == 0) { - writeEmpty(fieldInfo, hnswMeta); + int size = vectors.size(); + if (writeTrivialField(fieldInfo, size)) { return; } - try { int dimensions = fieldInfo.getVectorDimension(); - int bytesPerVector = (dimensions + 7) / 8; - - CuVSMatrix dataset = Utils.createByteMatrix(vectors, bytesPerVector); + int bytesPerVector = Math.ceilDiv(dimensions, 8); + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(size, bytesPerVector, CuVSMatrix.DataType.BYTE)) { + for (byte[] vector : vectors) { + builder.addVector(vector); + } + writeFieldInternal(fieldInfo, builder.build()); + } + } catch (Throwable t) { + throw Utils.handleThrowable(t); + } + } - if (dataset.size() < 2) { - writeSingleVectorGraph(fieldInfo, vectors); + /** Builds and writes an index from an owned binary-vector matrix. */ + private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { + try (Utils.OwnedIndex ownedIndex = Utils.ownDataset(dataset)) { + int size = Math.toIntExact(dataset.size()); + if (writeTrivialField(fieldInfo, size)) { return; } - + int dimensions = fieldInfo.getVectorDimension(); CagraIndexParams params = CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); - CagraIndex cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()) .withDataset(dataset) .withIndexParams(params) .build(); + ownedIndex.transferTo(cagraIndex); CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); - int size = (int) dataset.size(); - - // Create multi-layer HNSW graph from CAGRA GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( - fieldInfo, dimensions, adjacencyListMatrix, dataset, @@ -186,11 +193,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw QuantizationType.BINARY); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - // Write the graph to the vector index int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - - // Write metadata writeMeta( hnswVectorIndex, hnswMeta, @@ -200,14 +204,24 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw size, hnswGraph, graphLevelNodeOffsets); - - cagraIndex.close(); - } catch (Throwable t) { - Utils.handleThrowable(t); + throw Utils.handleThrowable(t); } } + /** Writes the empty or one-vector representation, if {@code size} is trivial. */ + private boolean writeTrivialField(FieldInfo fieldInfo, int size) throws IOException { + if (size == 0) { + writeEmpty(fieldInfo, hnswMeta); + return true; + } + if (size == 1) { + writeSingleVectorGraph(fieldInfo); + return true; + } + return false; + } + /** * Build the indexes and writes it to the disk. */ @@ -259,11 +273,9 @@ private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) thr * Builds and writes a single vector graph. * * @param fieldInfo instance of FieldInfo - * @param vectors the list of binary quantized vectors * @throws IOException I/O Exceptions */ - private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) - throws IOException { + private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { // Workaround for CAGRA not supporting single vector indexes try { int size = 1; @@ -310,20 +322,60 @@ public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOE */ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { try { - FloatVectorValues mergedVectorValues = - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + int dimensions = fieldInfo.getVectorDimension(); + BinaryQuantizer quantizer = collectBinaryMergeStats(fieldInfo, mergeState, dimensions); + if (writeTrivialField(fieldInfo, quantizer.count())) { + return; + } - if (mergedVectorValues != null) { - List floatVectors = new ArrayList<>(); - KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); - for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { - floatVectors.add(mergedVectorValues.vectorValue(iter.index()).clone()); + int bytesPerVector = Math.ceilDiv(dimensions, 8); + FloatVectorValues mergedVectors = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(quantizer.count(), bytesPerVector, CuVSMatrix.DataType.BYTE)) { + byte[] quantized = new byte[bytesPerVector]; + int encodedCount = 0; + KnnVectorValues.DocIndexIterator iterator = mergedVectors.iterator(); + for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) { + if (encodedCount >= quantizer.count()) { + throw new IOException( + "Merged vector count changed between quantization passes: expected " + + quantizer.count() + + ", received more rows"); + } + quantizer.quantize(mergedVectors.vectorValue(iterator.index()), quantized); + builder.addVector(quantized); + encodedCount = Math.incrementExact(encodedCount); + } + if (encodedCount != quantizer.count()) { + throw new IOException( + "Merged vector count changed between quantization passes: expected " + + quantizer.count() + + ", received " + + encodedCount); } - writeFieldInternal(fieldInfo, quantizeFloatVectorsToBinary(floatVectors)); + writeFieldInternal(fieldInfo, builder.build()); } } catch (Throwable t) { - Utils.handleThrowable(t); + throw Utils.handleThrowable(t); + } + } + + private static BinaryQuantizer collectBinaryMergeStats( + FieldInfo fieldInfo, MergeState mergeState, int dimensions) throws IOException { + FloatVectorValues mergedVectors = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + BinaryQuantizer quantizer = new BinaryQuantizer(dimensions); + if (mergedVectors != null) { + KnnVectorValues.DocIndexIterator iterator = mergedVectors.iterator(); + for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) { + quantizer.add(mergedVectors.vectorValue(iterator.index())); + } + } + if (quantizer.count() > 0) { + quantizer.finish(); } + return quantizer; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index f9dbbfb06d..0ecb660663 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -7,7 +7,6 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createMultiLayerHnswGraph; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createSingleVectorHnswGraph; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.printInfoStream; -import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToScalar; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeEmpty; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeGraph; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeMeta; @@ -25,6 +24,7 @@ import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; +import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.ScalarQuantizer; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -150,14 +150,6 @@ private static byte signedToUnsignedByte(byte signedByte) { return (byte) (signedByte & 0xFF); } - private static byte[] convertSignedToUnsigned(byte[] signedVector) { - byte[] unsignedVector = new byte[signedVector.length]; - for (int i = 0; i < signedVector.length; i++) { - unsignedVector[i] = signedToUnsignedByte(signedVector[i]); - } - return unsignedVector; - } - /** * Builds the intermediate CAGRA index and builds and writes the HNSW index. * @@ -166,43 +158,50 @@ private static byte[] convertSignedToUnsigned(byte[] signedVector) { * @throws IOException */ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { - if (vectors.size() == 0) { - writeEmpty(fieldInfo, hnswMeta); + int size = vectors.size(); + if (writeTrivialField(fieldInfo, size)) { return; } - try { int dimensions = fieldInfo.getVectorDimension(); - - // Convert 7-bit signed bytes to 8-bit unsigned bytes for cuVS compatibility - List unsignedVectors = new ArrayList<>(vectors.size()); - for (Object signedVector : vectors) { - unsignedVectors.add(convertSignedToUnsigned((byte[]) signedVector)); + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(size, dimensions, CuVSMatrix.DataType.BYTE)) { + byte[] unsignedVector = new byte[dimensions]; + for (Object value : vectors) { + if (!(value instanceof byte[] signedVector) || signedVector.length != dimensions) { + throw new IllegalArgumentException( + "Expected scalar-quantized byte[" + dimensions + "] vector"); + } + copySignedToUnsigned(signedVector, unsignedVector); + builder.addVector(unsignedVector); + } + writeFieldInternal(fieldInfo, builder.build()); } + } catch (Throwable t) { + throw Utils.handleThrowable(t); + } + } - // Create CuVSMatrix with BYTE data type (unsigned bytes) - CuVSMatrix dataset = Utils.createByteMatrix(unsignedVectors, dimensions); - - if (dataset.size() < 2) { - writeSingleVectorGraph(fieldInfo, unsignedVectors); + /** Builds and writes an index from an owned scalar-vector matrix. */ + private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { + try (Utils.OwnedIndex ownedIndex = Utils.ownDataset(dataset)) { + int size = Math.toIntExact(dataset.size()); + if (writeTrivialField(fieldInfo, size)) { return; } - + int dimensions = fieldInfo.getVectorDimension(); CagraIndexParams params = CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); - CagraIndex cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()) .withDataset(dataset) .withIndexParams(params) .build(); + ownedIndex.transferTo(cagraIndex); CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); - - int size = (int) dataset.size(); GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( - fieldInfo, dimensions, adjacencyListMatrix, dataset, @@ -211,13 +210,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE QuantizationType.SCALAR); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - - // Write the graph to the vector index int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); - long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - - // Write metadata writeMeta( hnswVectorIndex, hnswMeta, @@ -227,13 +221,24 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE size, hnswGraph, graphLevelNodeOffsets); - - cagraIndex.close(); } catch (Throwable t) { - Utils.handleThrowable(t); + throw Utils.handleThrowable(t); } } + /** Writes the empty or one-vector representation, if {@code size} is trivial. */ + private boolean writeTrivialField(FieldInfo fieldInfo, int size) throws IOException { + if (size == 0) { + writeEmpty(fieldInfo, hnswMeta); + return true; + } + if (size == 1) { + writeSingleVectorGraph(fieldInfo); + return true; + } + return false; + } + /** * Build the indexes and writes it to the disk. */ @@ -285,11 +290,9 @@ private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) thr * Builds and writes a single vector graph. * * @param fieldInfo instance of FieldInfo - * @param vectors the list of scalar quantized vectors (already converted to unsigned) * @throws IOException I/O Exceptions */ - private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) - throws IOException { + private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { // Workaround for CAGRA not supporting single vector indexes try { int size = 1; @@ -334,19 +337,65 @@ public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOE */ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { try { - FloatVectorValues mergedVectorValues = - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + int dimensions = fieldInfo.getVectorDimension(); + ScalarQuantizer quantizer = collectScalarMergeStats(fieldInfo, mergeState, dimensions); + if (writeTrivialField(fieldInfo, quantizer.count())) { + return; + } - if (mergedVectorValues != null) { - List floatVectors = new ArrayList<>(); - KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); - for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { - floatVectors.add(mergedVectorValues.vectorValue(iter.index()).clone()); + FloatVectorValues mergedVectors = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(quantizer.count(), dimensions, CuVSMatrix.DataType.BYTE)) { + byte[] quantized = new byte[dimensions]; + int encodedCount = 0; + KnnVectorValues.DocIndexIterator iterator = mergedVectors.iterator(); + for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) { + if (encodedCount >= quantizer.count()) { + throw new IOException( + "Merged vector count changed between quantization passes: expected " + + quantizer.count() + + ", received more rows"); + } + quantizer.quantize(mergedVectors.vectorValue(iterator.index()), quantized); + copySignedToUnsigned(quantized, quantized); + builder.addVector(quantized); + encodedCount = Math.incrementExact(encodedCount); } - writeFieldInternal(fieldInfo, quantizeFloatVectorsToScalar(floatVectors)); + if (encodedCount != quantizer.count()) { + throw new IOException( + "Merged vector count changed between quantization passes: expected " + + quantizer.count() + + ", received " + + encodedCount); + } + writeFieldInternal(fieldInfo, builder.build()); } } catch (Throwable t) { - Utils.handleThrowable(t); + throw Utils.handleThrowable(t); + } + } + + private static ScalarQuantizer collectScalarMergeStats( + FieldInfo fieldInfo, MergeState mergeState, int dimensions) throws IOException { + FloatVectorValues mergedVectors = + KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); + ScalarQuantizer quantizer = new ScalarQuantizer(dimensions); + if (mergedVectors != null) { + KnnVectorValues.DocIndexIterator iterator = mergedVectors.iterator(); + for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) { + quantizer.add(mergedVectors.vectorValue(iterator.index())); + } + } + if (quantizer.count() > 0) { + quantizer.finish(); + } + return quantizer; + } + + private static void copySignedToUnsigned(byte[] source, byte[] destination) { + for (int i = 0; i < source.length; i++) { + destination[i] = signedToUnsignedByte(source[i]); } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java index 9b8028543c..a47f65658c 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -61,12 +61,13 @@ static RuntimeException handleThrowable(Throwable t) throws IOException { * @return a host-memory CuVSMatrix */ static CuVSMatrix createFloatMatrix(List data, int dimensions) { - CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT); - for (float[] vector : data) { - builder.addVector(vector); + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT)) { + for (float[] vector : data) { + builder.addVector(vector); + } + return builder.build(); } - return builder.build(); } /** @@ -77,12 +78,13 @@ static CuVSMatrix createFloatMatrix(List data, int dimensions) { * @return a host-memory CuVSMatrix with BYTE data type */ static CuVSMatrix createByteMatrix(List data, int bytesPerVector) { - CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE); - for (byte[] vector : data) { - builder.addVector(vector); + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE)) { + for (byte[] vector : data) { + builder.addVector(vector); + } + return builder.build(); } - return builder.build(); } /** @@ -93,12 +95,93 @@ static CuVSMatrix createByteMatrix(List data, int bytesPerVector) { * @return a host-memory CuVSMatrix with BYTE data type */ static CuVSMatrix createByteMatrixFromArray(byte[][] data, int bytesPerVector) { - CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE); - for (byte[] vector : data) { - builder.addVector(vector); + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE)) { + for (byte[] vector : data) { + builder.addVector(vector); + } + return builder.build(); + } + } + + /** + * Closes an index that owns {@code dataset}. If index cleanup fails before releasing the + * dataset, a direct dataset close is attempted and attached to the index failure when needed. + */ + static void closeIndexWithDatasetFallback(AutoCloseable index, AutoCloseable dataset) + throws Exception { + try { + index.close(); + } catch (Throwable indexCloseFailure) { + try { + dataset.close(); + } catch (Throwable datasetCloseFailure) { + if (indexCloseFailure != datasetCloseFailure) { + indexCloseFailure.addSuppressed(datasetCloseFailure); + } + } + rethrowCloseFailure(indexCloseFailure); + } + } + + /** Starts an ownership scope for a dataset that may later be transferred to an index. */ + static OwnedIndex ownDataset(AutoCloseable dataset) { + return new OwnedIndex<>(dataset); + } + + /** + * Owns a dataset immediately and, after {@link #transferTo}, closes the owning index with a + * direct dataset-close fallback. + */ + static final class OwnedIndex implements AutoCloseable { + private AutoCloseable dataset; + private I index; + private boolean closed; + + private OwnedIndex(AutoCloseable dataset) { + this.dataset = java.util.Objects.requireNonNull(dataset, "dataset"); + } + + void transferTo(I index) { + if (closed || this.index != null) { + throw new IllegalStateException("Dataset ownership has already been transferred"); + } + this.index = java.util.Objects.requireNonNull(index, "index"); + } + + I index() { + if (index == null) { + throw new IllegalStateException("Dataset ownership has not been transferred to an index"); + } + return index; + } + + @Override + public void close() throws Exception { + if (closed) { + return; + } + closed = true; + AutoCloseable ownedDataset = dataset; + I ownedIndex = index; + dataset = null; + index = null; + if (ownedIndex == null) { + ownedDataset.close(); + } else { + closeIndexWithDatasetFallback(ownedIndex, ownedDataset); + } + } + } + + private static void rethrowCloseFailure(Throwable failure) throws Exception { + if (failure instanceof Exception exception) { + throw exception; + } + if (failure instanceof Error error) { + throw error; } - return builder.build(); + throw new AssertionError("Unexpected throwable from AutoCloseable.close()", failure); } /** diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java index cd4aaa6ec8..eef332bd85 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java @@ -11,7 +11,9 @@ import java.io.IOException; import java.util.ArrayList; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; import java.util.Random; import java.util.Set; import java.util.logging.Level; @@ -318,6 +320,64 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio } } + @Test + public void testForceMergeCountsOnlyLiveSparseVectors() throws IOException { + final String vectorField = "vector"; + final int dimensions = 129; + Map expected = new LinkedHashMap<>(); + + try (Directory directory = newDirectory()) { + try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + for (int segment = 0; segment < 3; segment++) { + for (int row = 0; row < 5; row++) { + String id = segment + "-" + row; + Document document = new Document(); + document.add(new StringField("id", id, Field.Store.YES)); + if (row < 4) { + float[] vector = deterministicVector(segment * 5 + row, dimensions); + document.add( + new KnnFloatVectorField(vectorField, vector, VectorSimilarityFunction.EUCLIDEAN)); + if (row != 1) { + expected.put(id, vector); + } + } + writer.addDocument(document); + } + writer.commit(); + } + for (int segment = 0; segment < 3; segment++) { + writer.deleteDocuments(new Term("id", segment + "-1")); + } + writer.commit(); + writer.forceMerge(1); + } + + TestUtil.checkIndex(directory); + try (DirectoryReader reader = DirectoryReader.open(directory)) { + var leaf = getOnlyLeafReader(reader); + var values = leaf.getFloatVectorValues(vectorField); + assertNotNull(values); + assertEquals(expected.size(), values.size()); + Set seen = new HashSet<>(); + for (int ordinal = 0; ordinal < values.size(); ordinal++) { + String id = leaf.storedFields().document(values.ordToDoc(ordinal)).get("id"); + assertTrue("Unexpected or duplicate vector for " + id, seen.add(id)); + assertNotNull("No expected vector for " + id, expected.get(id)); + assertArrayEquals(expected.get(id), values.vectorValue(ordinal), 0.0f); + } + assertEquals(expected.keySet(), seen); + } + } + } + + private static float[] deterministicVector(int id, int dimensions) { + float[] vector = new float[dimensions]; + for (int dimension = 0; dimension < dimensions; dimension++) { + vector[dimension] = id * 10.0f + dimension * 0.01f; + } + return vector; + } + private RandomIndexWriter createWriter(Directory directory) throws IOException { return new RandomIndexWriter( random(), diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java new file mode 100644 index 0000000000..16a5c8b84d --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java @@ -0,0 +1,171 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import com.carrotsearch.randomizedtesting.annotations.Name; +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; +import java.util.List; +import java.util.Random; +import java.util.Set; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.HnswGraphProvider; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.hnsw.HnswGraph; + +/** Exercises the native upper-layer path through serialization and CPU search. */ +@SuppressSysoutChecks(bugUrl = "") +public class TestAcceleratedHNSWMultiLayerRoundTrip extends LuceneTestCase { + + private static final String FIELD = "vector"; + private static final int VECTOR_COUNT = 256; + private static final int GRAPH_DEGREE = 16; + + private final KnnVectorsFormat format; + private final int dimensions; + + public TestAcceleratedHNSWMultiLayerRoundTrip( + @Name("knnVectorsFormat") KnnVectorsFormat format, @Name("dimensions") int dimensions) { + this.format = format; + this.dimensions = dimensions; + } + + @ParametersFactory + public static List parameters() { + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withIntermediateGraphDegree(32) + .withGraphDegree(GRAPH_DEGREE) + .withHNSWLayer(3) + .build(); + return Arrays.asList( + new Object[][] { + {new Lucene99AcceleratedHNSWVectorsFormat(params), 32}, + {new LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat(params), 129}, + {new LuceneAcceleratedHNSWScalarQuantizedVectorsFormat(params), 32} + }); + } + + public void testThreeLevelGraphCanBeReopenedAndSearched() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + float[][] vectors = randomVectors(VECTOR_COUNT, dimensions); + + try (Directory directory = newDirectory()) { + IndexWriterConfig config = + newIndexWriterConfig().setCodec(TestUtil.alwaysKnnVectorsFormat(format)); + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int id = 0; id < vectors.length; id++) { + Document document = new Document(); + document.add(new StringField("id", Integer.toString(id), Field.Store.YES)); + document.add(new KnnFloatVectorField(FIELD, vectors[id], EUCLIDEAN)); + writer.addDocument(document); + } + writer.forceMerge(1); + } + + TestUtil.checkIndex(directory); + try (DirectoryReader reader = DirectoryReader.open(directory)) { + LeafReader leaf = getOnlyLeafReader(reader); + HnswGraph graph = graphOf(leaf); + assertEquals(3, graph.numLevels()); + + List> nodesByLevel = collectNodes(graph); + int expectedLevelOneSize = Math.max(2, VECTOR_COUNT / graph.maxConn()); + int expectedLevelTwoSize = Math.max(2, expectedLevelOneSize / graph.maxConn()); + assertEquals(VECTOR_COUNT, nodesByLevel.get(0).size()); + assertEquals(expectedLevelOneSize, nodesByLevel.get(1).size()); + assertEquals(expectedLevelTwoSize, nodesByLevel.get(2).size()); + assertTrue(nodesByLevel.get(0).containsAll(nodesByLevel.get(1))); + assertTrue(nodesByLevel.get(1).containsAll(nodesByLevel.get(2))); + assertUpperNeighborsStayOnTheirLevel(graph, nodesByLevel); + + IndexSearcher searcher = new IndexSearcher(reader); + var hits = searcher.search(new KnnFloatVectorQuery(FIELD, vectors[0], 10), 10); + assertEquals(10, hits.scoreDocs.length); + boolean foundExactVector = false; + for (var hit : hits.scoreDocs) { + foundExactVector |= "0".equals(searcher.storedFields().document(hit.doc).get("id")); + } + assertTrue("the indexed vector must be returned for its own query", foundExactVector); + } + } + } + + private static List> collectNodes(HnswGraph graph) throws Exception { + List> nodesByLevel = new ArrayList<>(); + for (int level = 0; level < graph.numLevels(); level++) { + Set nodes = new HashSet<>(); + HnswGraph.NodesIterator iterator = graph.getNodesOnLevel(level); + while (iterator.hasNext()) { + assertTrue("duplicate node on level " + level, nodes.add(iterator.nextInt())); + } + nodesByLevel.add(nodes); + } + return nodesByLevel; + } + + private static void assertUpperNeighborsStayOnTheirLevel( + HnswGraph graph, List> nodesByLevel) throws Exception { + for (int level = 1; level < graph.numLevels(); level++) { + Set nodes = nodesByLevel.get(level); + int arcCount = 0; + for (int node : nodes) { + graph.seek(level, node); + for (int neighbor = graph.nextNeighbor(); + neighbor != NO_MORE_DOCS; + neighbor = graph.nextNeighbor()) { + assertTrue("negative upper-layer neighbor", neighbor >= 0); + assertTrue("upper-layer neighbor outside the index", neighbor < VECTOR_COUNT); + assertTrue( + "upper-layer neighbor is not a member of level " + level, nodes.contains(neighbor)); + arcCount++; + } + } + assertTrue("upper layer " + level + " contains no arcs", arcCount > 0); + } + } + + private static HnswGraph graphOf(LeafReader leaf) throws Exception { + KnnVectorsReader reader = ((CodecReader) leaf).getVectorReader(); + if (reader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + reader = fieldsReader.getFieldReader(FIELD); + } + return ((HnswGraphProvider) reader).getGraph(FIELD); + } + + private static float[][] randomVectors(int count, int dimensions) { + Random random = new Random(0x2476L); + float[][] vectors = new float[count][dimensions]; + for (float[] vector : vectors) { + for (int dimension = 0; dimension < vector.length; dimension++) { + vector[dimension] = random.nextFloat(); + } + } + return vectors; + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWQuantizers.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWQuantizers.java new file mode 100644 index 0000000000..52adb537c3 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWQuantizers.java @@ -0,0 +1,55 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.BinaryQuantizer; +import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.ScalarQuantizer; +import java.util.List; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.Test; + +public class TestAcceleratedHNSWQuantizers extends LuceneTestCase { + + @Test + public void testBinaryGoldenBytesAndReusableBufferClearing() { + float[] low = new float[10]; + float[] high = new float[] {2, 0, 2, 0, 2, 0, 2, 0, 2, 0}; + BinaryQuantizer quantizer = new BinaryQuantizer(10); + quantizer.add(low); + quantizer.add(high); + quantizer.finish(); + + byte[] scratch = new byte[] {(byte) 0xff, (byte) 0xff}; + quantizer.quantize(high, scratch); + assertArrayEquals(new byte[] {0x55, 0x01}, scratch); + + quantizer.quantize(low, scratch); + assertArrayEquals(new byte[] {0x00, 0x00}, scratch); + List quantized = AcceleratedHNSWUtils.quantizeFloatVectorsToBinary(List.of(low, high)); + assertArrayEquals(new byte[] {0x00, 0x00}, quantized.get(0)); + assertArrayEquals(new byte[] {0x55, 0x01}, quantized.get(1)); + } + + @Test + public void testScalarGoldenBytesPreserveExistingNegativeAndConstantBehavior() { + float[] first = new float[] {-4, -2, 7}; + float[] second = new float[] {-2, -2, 7}; + ScalarQuantizer quantizer = new ScalarQuantizer(3); + quantizer.add(first); + quantizer.add(second); + quantizer.finish(); + + byte[] scratch = new byte[3]; + quantizer.quantize(first, scratch); + assertArrayEquals(new byte[] {-64, -64, 0}, scratch); + quantizer.quantize(second, scratch); + assertArrayEquals(new byte[] {0, -64, 0}, scratch); + + List quantized = + AcceleratedHNSWUtils.quantizeFloatVectorsToScalar(List.of(first, second)); + assertArrayEquals(new byte[] {-64, -64, 0}, quantized.get(0)); + assertArrayEquals(new byte[] {0, -64, 0}, quantized.get(1)); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUpperLayers.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUpperLayers.java new file mode 100644 index 0000000000..aa20071797 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUpperLayers.java @@ -0,0 +1,273 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.RowView; +import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; +import java.io.IOException; +import java.lang.reflect.Method; +import java.lang.reflect.Modifier; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; +import org.junit.Test; + +public class TestAcceleratedHNSWUpperLayers extends LuceneTestCase { + + @Test + public void testLegacyListOverloadDescriptorIsPresent() throws Exception { + Method method = + AcceleratedHNSWUtils.class.getMethod( + "createMultiLayerHnswGraph", + FieldInfo.class, + int.class, + int.class, + CuVSMatrix.class, + List.class, + int.class, + CagraIndexParams.class, + QuantizationType.class); + + assertEquals(GPUBuiltHnswGraph.class, method.getReturnType()); + + Method matrixOverload = + AcceleratedHNSWUtils.class.getDeclaredMethod( + "createMultiLayerHnswGraph", + int.class, + CuVSMatrix.class, + CuVSMatrix.class, + int.class, + CagraIndexParams.class, + QuantizationType.class); + assertFalse(Modifier.isPublic(matrixOverload.getModifiers())); + } + + @Test + public void testRemapPreservesSentinelAndMapsAbsoluteOrdinals() throws IOException { + int[] destination = new int[3]; + + AcceleratedHNSWUtils.remapSubsetAdjacencyRow( + 0, new IntRow(-1, 2, 0), 3, new int[] {2, 7, 11}, destination); + + assertArrayEquals(new int[] {-1, 11, 2}, destination); + } + + @Test + public void testRemapRejectsPositiveOrdinalOutsideSubset() { + IOException thrown = + assertThrows( + IOException.class, + () -> + AcceleratedHNSWUtils.remapSubsetAdjacencyRow( + 4, new IntRow(0, 3), 2, new int[] {2, 7, 11}, new int[2])); + + assertTrue(thrown.getMessage().contains("row 4, column 1")); + assertTrue(thrown.getMessage().contains("outside [0, 3)")); + } + + @Test + public void testGraphMaterializationSkipsNativeSentinels() { + IntMatrix adjacency = new IntMatrix(new int[][] {{-1, 1, 2}, {0, -1, 2}, {0, 1, -1}}); + List layerNodes = new ArrayList<>(); + layerNodes.add(null); + + GPUBuiltHnswGraph graph = new GPUBuiltHnswGraph(3, 2, layerNodes, List.of(adjacency)); + + assertArrayEquals( + new int[] {1, 2}, + Arrays.copyOf(graph.getNeighbors(0, 0).nodes(), graph.getNeighbors(0, 0).size())); + assertEquals(2, graph.getNeighbors(0, 0).size()); + } + + @Test + public void testGraphMaterializationRejectsPositiveOrdinalOutsideGraph() { + IntMatrix adjacency = new IntMatrix(new int[][] {{1}, {2}}); + List layerNodes = new ArrayList<>(); + layerNodes.add(null); + + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, + () -> new GPUBuiltHnswGraph(2, 2, layerNodes, List.of(adjacency))); + + assertTrue(thrown.getMessage().contains("ordinal 2 outside graph size 2")); + } + + @Test + public void testUpperLayerNeighborLookupUsesSortedNodeOrdinals() { + IntMatrix baseAdjacency = new IntMatrix(new int[][] {{-1}, {-1}, {-1}, {-1}, {-1}, {-1}, {-1}}); + IntMatrix upperAdjacency = new IntMatrix(new int[][] {{3}, {6}, {1}}); + int[] upperNodes = new int[] {1, 3, 6}; + GPUBuiltHnswGraph graph = + new GPUBuiltHnswGraph( + 7, 2, Arrays.asList((int[]) null, upperNodes), List.of(baseAdjacency, upperAdjacency)); + + Arrays.fill(upperNodes, 0); + assertArrayEquals(new int[] {1, 3, 6}, NodesIterator.getSortedNodes(graph.getNodesOnLevel(1))); + assertArrayEquals(new int[] {3}, neighbors(graph, 1, 1)); + assertArrayEquals(new int[] {6}, neighbors(graph, 1, 3)); + assertArrayEquals(new int[] {1}, neighbors(graph, 1, 6)); + assertNull(graph.getNeighbors(1, 0)); + assertNull(graph.getNeighbors(1, 4)); + assertNull(graph.getNeighbors(1, 7)); + } + + @Test + public void testUpperLayerNeighborLookupPreservesUnsortedPublicInput() { + IntMatrix baseAdjacency = new IntMatrix(new int[][] {{-1}, {-1}, {-1}, {-1}, {-1}, {-1}, {-1}}); + IntMatrix upperAdjacency = new IntMatrix(new int[][] {{1}, {3}, {6}}); + GPUBuiltHnswGraph graph = + new GPUBuiltHnswGraph( + 7, + 2, + Arrays.asList((int[]) null, new int[] {6, 1, 3}), + List.of(baseAdjacency, upperAdjacency)); + + assertArrayEquals(new int[] {1}, neighbors(graph, 1, 6)); + assertArrayEquals(new int[] {3}, neighbors(graph, 1, 1)); + assertArrayEquals(new int[] {6}, neighbors(graph, 1, 3)); + assertNull(graph.getNeighbors(1, 2)); + } + + @Test + public void testUpperLayerNodeValidationRejectsDuplicatesAndOutOfRangeOrdinals() { + IntMatrix baseAdjacency = new IntMatrix(new int[][] {{-1}, {-1}, {-1}, {-1}}); + IntMatrix threeRows = new IntMatrix(new int[][] {{1}, {2}, {3}}); + IllegalArgumentException duplicate = + assertThrows( + IllegalArgumentException.class, + () -> + new GPUBuiltHnswGraph( + 4, + 2, + Arrays.asList((int[]) null, new int[] {1, 2, 1}), + List.of(baseAdjacency, threeRows))); + assertTrue(duplicate.getMessage().contains("duplicate node ordinal 1")); + + IntMatrix twoRows = new IntMatrix(new int[][] {{1}, {2}}); + IllegalArgumentException outOfRange = + assertThrows( + IllegalArgumentException.class, + () -> + new GPUBuiltHnswGraph( + 4, + 2, + Arrays.asList((int[]) null, new int[] {1, 4}), + List.of(baseAdjacency, twoRows))); + assertTrue(outOfRange.getMessage().contains("ordinal 4 outside [0, 4)")); + } + + private static int[] neighbors(GPUBuiltHnswGraph graph, int level, int node) { + var neighbors = graph.getNeighbors(level, node); + return Arrays.copyOf(neighbors.nodes(), neighbors.size()); + } + + private record IntRow(int... values) implements RowView { + @Override + public long size() { + return values.length; + } + + @Override + public int getAsInt(long index) { + return values[Math.toIntExact(index)]; + } + + @Override + public float getAsFloat(long index) { + throw new UnsupportedOperationException(); + } + + @Override + public byte getAsByte(long index) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(int[] array) { + System.arraycopy(values, 0, array, 0, values.length); + } + + @Override + public void toArray(float[] array) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(byte[] array) { + throw new UnsupportedOperationException(); + } + } + + private record IntMatrix(int[][] values) implements CuVSMatrix { + @Override + public long size() { + return values.length; + } + + @Override + public long columns() { + return values[0].length; + } + + @Override + public DataType dataType() { + return DataType.INT; + } + + @Override + public RowView getRow(long row) { + return new IntRow(values[Math.toIntExact(row)]); + } + + @Override + public void toArray(int[][] array) { + for (int row = 0; row < values.length; row++) { + System.arraycopy(values[row], 0, array[row], 0, values[row].length); + } + } + + @Override + public void toArray(float[][] array) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(byte[][] array) { + throw new UnsupportedOperationException(); + } + + @Override + public void toHost(CuVSHostMatrix hostMatrix) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSHostMatrix toHost() { + throw new UnsupportedOperationException(); + } + + @Override + public void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources resources) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSDeviceMatrix toDevice(CuVSResources resources) { + throw new UnsupportedOperationException(); + } + + @Override + public void close() {} + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsWriterFailureHandling.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsWriterFailureHandling.java new file mode 100644 index 0000000000..5062a56799 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsWriterFailureHandling.java @@ -0,0 +1,150 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.LibraryException; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.junit.Test; + +public class TestCuVS2510GPUVectorsWriterFailureHandling extends LuceneTestCase { + + @Test + public void testNativeConstructionFailureIsRecoverableBeforePersistence() { + LibraryException nativeFailure = new LibraryException("native construction"); + + Throwable classified = + CuVS2510GPUVectorsWriter.classifyCagraWriteFailure(nativeFailure, false, false); + + assertTrue( + classified instanceof CuVS2510GPUVectorsWriter.RecoverableCagraConstructionException); + assertSame(nativeFailure, classified.getCause()); + } + + @Test + public void testArbitraryConstructionFailureIsNotRecoverable() { + RuntimeException programmingFailure = new IllegalStateException("programming failure"); + + Throwable classified = + CuVS2510GPUVectorsWriter.classifyCagraWriteFailure(programmingFailure, false, false); + + assertSame(programmingFailure, classified); + } + + @Test + public void testNativeFailureIsNotRecoverableAfterPersistenceStarts() { + LibraryException nativeFailure = new LibraryException("serialization"); + + Throwable classified = + CuVS2510GPUVectorsWriter.classifyCagraWriteFailure(nativeFailure, true, false); + + assertSame(nativeFailure, classified); + } + + @Test + public void testCleanupFailureDisablesFallback() { + LibraryException nativeFailure = new LibraryException("native construction"); + nativeFailure.addSuppressed(new IOException("cleanup")); + + Throwable classified = + CuVS2510GPUVectorsWriter.classifyCagraWriteFailure(nativeFailure, false, true); + + assertSame(nativeFailure, classified); + } + + @Test + public void testSuppressedFailureDisablesFallback() { + LibraryException nativeFailure = new LibraryException("native construction"); + nativeFailure.addSuppressed(new IOException("cleanup")); + + Throwable classified = + CuVS2510GPUVectorsWriter.classifyCagraWriteFailure(nativeFailure, false, false); + + assertSame(nativeFailure, classified); + } + + @Test + public void testPersistenceBoundaryIsMonotonic() { + CuVS2510GPUVectorsWriter.CagraWriteContext context = + new CuVS2510GPUVectorsWriter.CagraWriteContext(); + + assertFalse(context.persistenceStarted()); + context.beginPersistence(); + assertTrue(context.persistenceStarted()); + } + + @Test + public void testChangedOutputPositionRejectsFallback() { + Throwable failure = new LibraryException("native construction"); + + IOException thrown = + assertThrows( + IOException.class, + () -> CuVS2510GPUVectorsWriter.ensureFallbackOutputUnchanged(10L, 11L, failure)); + + assertSame(failure, thrown.getCause()); + assertTrue(thrown.getMessage().contains("changed from 10 to 11")); + } + + @Test + public void testUnchangedOutputPositionAllowsFallback() throws IOException { + CuVS2510GPUVectorsWriter.ensureFallbackOutputUnchanged( + 10L, 10L, new LibraryException("native construction")); + } + + @Test + public void testCagraResourcesCloseInDependencyOrder() { + List closeOrder = new ArrayList<>(); + AtomicInteger directDatasetCloses = new AtomicInteger(); + + Throwable failure = + CuVS2510GPUVectorsWriter.closeCagraResources( + () -> closeOrder.add("index"), + directDatasetCloses::incrementAndGet, + () -> closeOrder.add("padded-dataset"), + () -> closeOrder.add("device-vectors")); + + assertNull(failure); + assertEquals(List.of("index", "padded-dataset", "device-vectors"), closeOrder); + assertEquals(0, directDatasetCloses.get()); + } + + @Test + public void testBodyAndCleanupFailuresPreserveOrderAndSuppression() { + List closeOrder = new ArrayList<>(); + IOException bodyFailure = new IOException("serialize"); + RuntimeException indexFailure = new RuntimeException("index close"); + RuntimeException datasetFailure = new RuntimeException("dataset close"); + RuntimeException paddedFailure = new RuntimeException("padded close"); + RuntimeException deviceFailure = new RuntimeException("device close"); + + Throwable cleanupFailure = + CuVS2510GPUVectorsWriter.closeCagraResources( + failingCloseable("index", closeOrder, indexFailure), + failingCloseable("dataset", closeOrder, datasetFailure), + failingCloseable("padded-dataset", closeOrder, paddedFailure), + failingCloseable("device-vectors", closeOrder, deviceFailure)); + Throwable combined = + CuVS2510GPUVectorsWriter.combineOperationAndCleanupFailures(bodyFailure, cleanupFailure); + + assertSame(bodyFailure, combined); + assertEquals(List.of("index", "dataset", "padded-dataset", "device-vectors"), closeOrder); + assertArrayEquals(new Throwable[] {indexFailure}, bodyFailure.getSuppressed()); + assertArrayEquals( + new Throwable[] {datasetFailure, paddedFailure, deviceFailure}, + indexFailure.getSuppressed()); + } + + private static AutoCloseable failingCloseable( + String name, List closeOrder, RuntimeException failure) { + return () -> { + closeOrder.add(name); + throw failure; + }; + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java index 3b05beaacf..df493c3e22 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java @@ -12,7 +12,11 @@ import com.carrotsearch.randomizedtesting.annotations.Name; import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; import java.util.Arrays; +import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; +import java.util.Set; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; @@ -26,6 +30,7 @@ import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.Term; import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.store.ByteBuffersDirectory; import org.apache.lucene.store.Directory; @@ -166,6 +171,122 @@ public void testCosineSimilarity() throws Exception { } } + public void testForceMergeUsesOnlyLiveSparseVectors() throws Exception { + final String vectorField = "vector"; + final int dimensions = 129; + Map expected = new LinkedHashMap<>(); + + try (Directory directory = newDirectory(new ByteBuffersDirectory())) { + try (IndexWriter writer = new IndexWriter(directory, newIndexWriterConfig())) { + for (int segment = 0; segment < 3; segment++) { + for (int row = 0; row < 5; row++) { + String id = segment + "-" + row; + Document document = new Document(); + document.add(new StringField("id", id, Field.Store.YES)); + if (row < 4) { + float[] vector = deterministicVector(segment * 5 + row, dimensions); + document.add(new KnnFloatVectorField(vectorField, vector, EUCLIDEAN)); + if (row != 1) { + expected.put(id, vector); + } + } + writer.addDocument(document); + } + writer.commit(); + } + for (int segment = 0; segment < 3; segment++) { + writer.deleteDocuments(new Term("id", segment + "-1")); + } + writer.commit(); + writer.forceMerge(1); + } + + TestUtil.checkIndex(directory); + try (DirectoryReader reader = DirectoryReader.open(directory)) { + LeafReader leaf = getOnlyLeafReader(reader); + FloatVectorValues values = leaf.getFloatVectorValues(vectorField); + assertNotNull(values); + assertEquals(expected.size(), values.size()); + + Set seen = new HashSet<>(); + for (int ordinal = 0; ordinal < values.size(); ordinal++) { + int documentId = values.ordToDoc(ordinal); + String id = leaf.storedFields().document(documentId).get("id"); + assertTrue("Unexpected or duplicate vector for " + id, seen.add(id)); + assertArrayEquals(expected.get(id), values.vectorValue(ordinal), 0.0f); + } + assertEquals(expected.keySet(), seen); + + for (Map.Entry entry : expected.entrySet()) { + var hits = + leaf.searchNearestVectors( + vectorField, entry.getValue(), expected.size(), null, 1_000); + boolean found = false; + for (var hit : hits.scoreDocs) { + found |= entry.getKey().equals(leaf.storedFields().document(hit.doc).get("id")); + } + assertTrue("Exact vector was not searchable for " + entry.getKey(), found); + } + } + } + } + + public void testForceMergeWithZeroLiveVectors() throws Exception { + assertTrivialLiveVectorMerge(0); + } + + public void testForceMergeWithOneLiveVector() throws Exception { + assertTrivialLiveVectorMerge(1); + } + + private void assertTrivialLiveVectorMerge(int liveVectors) throws Exception { + final String vectorField = "vector"; + final int dimensions = 129; + try (Directory directory = newDirectory(new ByteBuffersDirectory())) { + try (IndexWriter writer = new IndexWriter(directory, newIndexWriterConfig())) { + for (int id = 0; id < 3; id++) { + Document vectorDocument = new Document(); + vectorDocument.add(new StringField("id", "vector-" + id, Field.Store.YES)); + vectorDocument.add( + new KnnFloatVectorField(vectorField, deterministicVector(id, dimensions), EUCLIDEAN)); + writer.addDocument(vectorDocument); + + Document sparseDocument = new Document(); + sparseDocument.add(new StringField("id", "sparse-" + id, Field.Store.YES)); + writer.addDocument(sparseDocument); + writer.commit(); + } + for (int id = liveVectors; id < 3; id++) { + writer.deleteDocuments(new Term("id", "vector-" + id)); + } + writer.commit(); + writer.forceMerge(1); + } + + TestUtil.checkIndex(directory); + try (DirectoryReader reader = DirectoryReader.open(directory)) { + LeafReader leaf = getOnlyLeafReader(reader); + FloatVectorValues values = leaf.getFloatVectorValues(vectorField); + if (liveVectors == 0) { + assertTrue(values == null || values.size() == 0); + } else { + assertNotNull(values); + assertEquals(1, values.size()); + assertEquals("vector-0", leaf.storedFields().document(values.ordToDoc(0)).get("id")); + assertArrayEquals(deterministicVector(0, dimensions), values.vectorValue(0), 0.0f); + } + } + } + } + + private static float[] deterministicVector(int id, int dimensions) { + float[] vector = new float[dimensions]; + for (int dimension = 0; dimension < dimensions; dimension++) { + vector[dimension] = id * 10.0f + dimension * 0.01f; + } + return vector; + } + @Override protected VectorEncoding randomVectorEncoding() { return VectorEncoding.FLOAT32; diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java index 6aead1bd8d..6fdfa5a2a0 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java @@ -47,4 +47,73 @@ public void testHandleThrowableWrapsCheckedExceptionWithCause() { assertSame(exception, thrown.getCause()); } + + @Test + public void testOwnedIndexClosesUntransferredDatasetOnce() throws Exception { + TrackingCloseable dataset = new TrackingCloseable(null); + Utils.OwnedIndex owned = Utils.ownDataset(dataset); + + owned.close(); + owned.close(); + + assertEquals(1, dataset.closeCount); + } + + @Test + public void testOwnedIndexDoesNotDirectlyCloseDatasetAfterSuccessfulIndexClose() + throws Exception { + TrackingCloseable dataset = new TrackingCloseable(null); + TrackingCloseable index = new TrackingCloseable(null); + Utils.OwnedIndex owned = Utils.ownDataset(dataset); + owned.transferTo(index); + + owned.close(); + + assertEquals(1, index.closeCount); + assertEquals(0, dataset.closeCount); + } + + @Test + public void testOwnedIndexPreservesBodyAndNestedCleanupFailures() { + IOException bodyFailure = new IOException("body"); + IOException indexCloseFailure = new IOException("index close"); + IOException datasetCloseFailure = new IOException("dataset close"); + TrackingCloseable dataset = new TrackingCloseable(datasetCloseFailure); + TrackingCloseable index = new TrackingCloseable(indexCloseFailure); + + IOException thrown = + assertThrows( + IOException.class, + () -> { + try (Utils.OwnedIndex owned = Utils.ownDataset(dataset)) { + owned.transferTo(index); + throw bodyFailure; + } + }); + + assertSame(bodyFailure, thrown); + assertEquals(1, thrown.getSuppressed().length); + assertSame(indexCloseFailure, thrown.getSuppressed()[0]); + assertEquals(1, indexCloseFailure.getSuppressed().length); + assertSame(datasetCloseFailure, indexCloseFailure.getSuppressed()[0]); + assertEquals(1, index.closeCount); + assertEquals(1, dataset.closeCount); + } + + private static final class TrackingCloseable implements AutoCloseable { + private final Exception failure; + private int closeCount; + + private TrackingCloseable(Exception failure) { + this.failure = failure; + } + + @Override + public void close() throws Exception { + closeCount++; + if (failure != null) { + throw failure; + } + } + } } From 1e1fe2f995cdd942676317e5c3e71ce41fbf95f4 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 18 Sep 2026 13:01:31 +0000 Subject: [PATCH 05/21] Harden device matrix construction and cleanup Acquire device streams before allocation, preserve ownership across builder cleanup, and keep CAGRA persistence on the original device matrix. Expand lifecycle, fallback, merge, and graph-integrity regression coverage. --- .../java-api-com-nvidia-cuvs-cuvsmatrix.md | 58 +++-- ...nvidia-cuvs-lucene-acceleratedhnswutils.md | 36 ++- ...ia-cuvs-lucene-cuvs2510gpuvectorswriter.md | 45 +++- ...om-nvidia-cuvs-lucene-gpubuilthnswgraph.md | 24 +- ...ne-lucene99acceleratedhnswvectorswriter.md | 16 +- ...leratedhnswbinaryquantizedvectorswriter.md | 10 +- ...leratedhnswscalarquantizedvectorswriter.md | 10 +- ...lucene-api-com-nvidia-cuvs-lucene-utils.md | 67 +++--- fern/pages/other/multidimensional_arrays.md | 12 +- .../com/nvidia/cuvs/CuVSDeviceMatrix.java | 6 +- .../java/com/nvidia/cuvs/CuVSHostMatrix.java | 7 +- .../nvidia/cuvs/MatrixBuilderLifecycle.java | 47 ++++ .../com/nvidia/cuvs/spi/JDKProvider.java | 86 ++++++-- .../nvidia/cuvs/MatrixBuilderLifecycleIT.java | 137 ++++++++++++ .../com/nvidia/cuvs/spi/CuVSProviderIT.java | 71 ++++++ .../cuvs/lucene/AcceleratedHNSWUtils.java | 104 ++++----- .../cuvs/lucene/CuVS2510GPUVectorsWriter.java | 91 ++++++-- .../cuvs/lucene/MatrixBuilderLifecycle.java | 64 ++++++ .../java/com/nvidia/cuvs/lucene/Utils.java | 52 +++-- .../TestAcceleratedHNSWDeletedDocuments.java | 138 +++++++++++- ...estAcceleratedHNSWMultiLayerRoundTrip.java | 11 +- .../TestCagraIndexAtAlignedDimensions.java | 52 ++++- ...VS2510GPUVectorsWriterFailureHandling.java | 205 +++++++++++++++++- .../lucene/TestQuantizedVectorsFormats.java | 36 ++- .../lucene/TestUtilsThrowableHandling.java | 112 ++++++++++ 25 files changed, 1255 insertions(+), 242 deletions(-) create mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/MatrixBuilderLifecycle.java create mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/MatrixBuilderLifecycleIT.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MatrixBuilderLifecycle.java diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsmatrix.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsmatrix.md index 29e3e618fd..0d62e838f4 100644 --- a/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsmatrix.md +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsmatrix.md @@ -114,6 +114,34 @@ Adds a single vector to the matrix. Each element is a raw float16 bit pattern st _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:102`_ +### build + +```java +T build() +``` + +Completes the matrix and transfers ownership to the caller. + +If this method fails, closing the builder releases any matrix storage allocated while the +builder was created. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:110`_ + +### close + +```java +@Override default void close() +``` + +Closes this builder. Built-in builders release matrix storage unless ownership was +transferred by a successful `#build()`. + +The default implementation preserves compatibility with providers compiled before +builders became closeable. Builders that allocate storage before `#build()` should +override this method. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:120`_ + ### hostBuilder ```java @@ -134,7 +162,7 @@ Returns a builder to create a new instance of a host-memory matrix a builder for creating a `CuVSHostMatrix` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:115`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:132`_ ### hostBuilder @@ -158,7 +186,7 @@ Returns a builder to create a new instance of a host-memory matrix a builder for creating a `CuVSDeviceMatrix` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:129`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:146`_ ### deviceBuilder @@ -181,7 +209,7 @@ Returns a builder to create a new instance of a dataset a builder for creating a `CuVSDeviceMatrix` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:144`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:161`_ ### deviceBuilder @@ -206,7 +234,7 @@ Returns a builder to create a new instance of a dataset a builder for creating a `CuVSDeviceMatrix` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:160`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:177`_ ### size @@ -220,7 +248,7 @@ Gets the size of the dataset Size of the dataset -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:176`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:193`_ ### columns @@ -235,7 +263,7 @@ or the graph degree for the graph represented as a list of neighbours Dimensions of the vectors in the dataset -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:184`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:201`_ ### dataType @@ -249,7 +277,7 @@ Gets the element type a `DataType` describing the matrix element type -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:191`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:208`_ ### getRow @@ -265,7 +293,7 @@ Get a view (0-copy) of the row data, as a list of integers (32 bit) | --- | --- | | `row` | the row for which to return the data | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:198`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:215`_ ### toArray @@ -281,7 +309,7 @@ Copies the content of this dataset to an on-heap Java matrix (array of arrays). | --- | --- | | `array` | the destination array. Must be of length `CuVSMatrix#size()` or bigger, and each element must be of length `CuVSMatrix#columns()` or bigger. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:206`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:223`_ ### toArray @@ -297,7 +325,7 @@ Copies the content of this dataset to an on-heap Java matrix (array of arrays). | --- | --- | | `array` | the destination array. Must be of length `CuVSMatrix#size()` or bigger, and each element must be of length `CuVSMatrix#columns()` or bigger. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:214`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:231`_ ### toArray @@ -313,7 +341,7 @@ Copies the content of this dataset to an on-heap Java matrix (array of arrays). | --- | --- | | `array` | the destination array. Must be of length `CuVSMatrix#size()` or bigger, and each element must be of length `CuVSMatrix#columns()` or bigger. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:222`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:239`_ ### toHost @@ -331,7 +359,7 @@ same element type and dimension. | --- | --- | | `hostMatrix` | the host-memory-backed matrix to fill. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:231`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:248`_ ### toHost @@ -345,7 +373,7 @@ the device matrix. The returned host matrix will need to be managed by the caller, which will be responsible to call `CuVSMatrix#close()` to free its resources when done. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:240`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:257`_ ### toDevice @@ -363,7 +391,7 @@ same element type and dimension. | --- | --- | | `deviceMatrix` | the device-memory-backed matrix to fill. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:249`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:266`_ ### toDevice @@ -377,6 +405,6 @@ the host matrix. The returned device matrix will need to be managed by the caller, which will be responsible to call `CuVSMatrix#close()` to free its resources when done. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:258`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:275`_ _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:17`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md index 219149bc6d..0aa315d910 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md @@ -29,13 +29,25 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWU public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, int size, int dimensions, CuVSMatrix adjacencyListMatrix, List vectors, int hnswLayers, CagraIndexParams params, QuantizationType quantization) throws Throwable ``` -Creates a multi-layer HNSW graph with dynamic number of layers. -M = ceil(cagraGraphDegree / 2), where cagraGraphDegree is the CAGRA adjacency list's degree -(its column count). Ceil is used to accommodate odd graph degrees. -Each layer contains 1/M nodes from the previous layer -Creates layers until the highest layer has ≤ M nodes +Creates a multi-layer HNSW graph from heap vectors. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:81`_ +This overload preserves the original public API. Only rows selected for an upper layer are +copied into native memory. + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:74`_ + +### createMultiLayerHnswGraph + +```java +static GPUBuiltHnswGraph createMultiLayerHnswGraph( int dimensions, CuVSMatrix adjacencyListMatrix, CuVSMatrix vectorDataset, int hnswLayers, CagraIndexParams params, QuantizationType quantization) throws Throwable +``` + +Creates a multi-layer HNSW graph from a native matrix. + +Only sampled rows are copied into each upper-layer matrix; the complete dataset is never +materialized on the Java heap. + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:106`_ ### writeGraph @@ -62,7 +74,7 @@ a 2D array of offsets | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:237`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:443`_ ### writeMeta @@ -91,7 +103,7 @@ Writes the meta information for the index. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:302`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:508`_ ### printInfoStream @@ -107,7 +119,7 @@ A utility method to print info/debugging messages using InfoStream. | --- | --- | | `msg` | the debugging message to print | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:384`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:590`_ ### writeEmpty @@ -129,7 +141,7 @@ Writes an empty meta information for the field. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:396`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:602`_ ### quantizeFloatVectorsToBinary @@ -152,7 +164,7 @@ Bits are packed: 8 dimensions per byte. A list of byte binary representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:409`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:615`_ ### quantizeFloatVectorsToScalar @@ -172,6 +184,6 @@ Scalar quantization. A list of byte scalar representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:451`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:643`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:31`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorswriter.md index df1abbbdc7..956074833f 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorswriter.md @@ -22,7 +22,7 @@ CAGRA(true, false), /** Builds a Brute Force index. */ BRUTE_FORCE(false, true), Builds a CAGRA index. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:89`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:91`_ ### BRUTE_FORCE @@ -32,7 +32,7 @@ BRUTE_FORCE(false, true), /** Builds both - CAGRA and Brute Force indexes. */ CA Builds a Brute Force index. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:92`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:94`_ ### CAGRA_AND_BRUTE_FORCE @@ -42,7 +42,7 @@ CAGRA_AND_BRUTE_FORCE(true, true) Builds both - CAGRA and Brute Force indexes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:95`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:97`_ ### CuVS2510GPUVectorsWriter @@ -66,7 +66,7 @@ Initializes `CuVS2510GPUVectorsWriter`. | --- | --- | | `IOException` | I/O exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:121`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:123`_ ### addField @@ -76,7 +76,30 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:160`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:162`_ + +### prepareCagraDataset + +```java +static CuVSDeviceMatrix prepareCagraDataset( CagraDatasetBuilderFactory builderFactory, CagraDatasetPopulation population) throws Throwable +``` + +Builds the device input and applies the pre-persistence fallback policy only after builder +cleanup succeeds. A builder cleanup failure is fatal; after ownership has transferred, this +method also closes the dataset before propagating that failure. A builder-factory failure is +also fatal because no cleanup handle was returned. + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:276`_ + +### closeCagraResources + +```java +static Throwable closeCagraResources( AutoCloseable index, AutoCloseable originalDataset, AutoCloseable indexDataset) +``` + +Closes CAGRA resources in dependency order and returns the first cleanup failure. + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:391`_ ### flush @@ -86,7 +109,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Creates the CAGRA and/or brute force indexes and writes them to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:303`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:471`_ ### mergeOneField @@ -96,7 +119,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:685`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:853`_ ### ramBytesUsed @@ -106,7 +129,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:696`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:864`_ ### finish @@ -116,7 +139,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:708`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:876`_ ### close @@ -126,6 +149,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Close the applicable resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:728`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:896`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:59`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:61`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md index f8a3aee9dc..91ed1f3a3d 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md @@ -31,7 +31,7 @@ Multi-layer constructor that supports arbitrary number of layers. | `layerNodes` | the nodes on the layer | | `layerAdjacencies` | adjacency list | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:41`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:44`_ ### getNodesOnLevel @@ -41,7 +41,7 @@ public NodesIterator getNodesOnLevel(int level) Get all nodes on a given level as node 0th ordinals. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:89`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:161`_ ### getNeighbors @@ -62,7 +62,7 @@ Get the neighbors for the node and the level it resides. an instance of NeighborArray -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:107`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:179`_ ### seek @@ -72,7 +72,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Move the pointer to exactly the given level's target. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:132`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:206`_ ### nextNeighbor @@ -82,7 +82,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Iterates over the neighbor list. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:142`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:216`_ ### entryNode @@ -92,7 +92,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns graph's entry point on the top level. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:173`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:247`_ ### maxConn @@ -102,7 +102,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap returns M, the maximum number of connections for a node. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:192`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:266`_ ### neighborCount @@ -112,7 +112,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns the neighbor count. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:207`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:281`_ ### size @@ -122,7 +122,7 @@ public int size() Returns the number of nodes in the graph. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:282`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:356`_ ### numLevels @@ -136,7 +136,7 @@ Returns the number of levels in the HNSW graph. the number of levels -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:291`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:365`_ ### dimensions @@ -150,6 +150,6 @@ Gets the vector dimension. the vector dimension -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:300`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:374`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:21`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:23`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md index 2f30caeb07..0bd4d8b8b3 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md @@ -37,7 +37,7 @@ Initializes `Lucene99AcceleratedHNSWVectorsWriter` | --- | --- | | `IOException` | IOException | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:86`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:89`_ ### addField @@ -47,7 +47,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:128`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:131`_ ### flush @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:203`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:226`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:291`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:339`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:300`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:348`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:320`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:368`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:330`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:378`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:52`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:55`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md index 91a4e67303..98b9ee0aa4 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:215`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:227`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:302`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:312`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:333`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:383`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:353`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:403`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:362`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:412`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:57`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md index 9c51cb26e9..04f5957faa 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:241`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:244`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:326`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:327`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:357`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:404`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:377`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:424`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:386`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:433`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:56`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md index 0d2a9cb1a6..be145bc37a 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md @@ -50,38 +50,34 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:43`_ ### createFloatMatrix ```java -static CuVSMatrix createFloatMatrix(List data, int dimensions, CuVSResources resources) +static CuVSMatrix createFloatMatrix(List data, int dimensions) throws IOException ``` -A method to build a CuVSMatrix from a list of float vectors. +Builds a host-memory CuVSMatrix from a list of float vectors. -Uses CuVSMatrix.Builder to copy vectors directly to device memory -without creating intermediate heap arrays. +Copies vectors directly into a native host matrix via `CuVSMatrix#hostBuilder`, +without creating an intermediate `float[][]` on the heap. **Parameters** | Name | Description | | --- | --- | | `data` | The float vectors | -| `dimensions` | The number float elements in each vector | -| `resources` | The CuVS resources for device matrix creation | +| `dimensions` | The number of float elements in each vector | **Returns** -an instance of CuVSMatrix +a host-memory CuVSMatrix -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:63`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:62`_ ### createByteMatrix ```java -static CuVSMatrix createByteMatrix( List data, int bytesPerVector, CuVSResources resources) +static CuVSMatrix createByteMatrix(List data, int bytesPerVector) throws IOException ``` -A method to build a CuVSMatrix from a list of byte vectors (for binary quantized vectors). - -Uses CuVSMatrix.Builder to copy vectors directly to device memory -without creating intermediate heap arrays. +Builds a host-memory CuVSMatrix from a list of byte vectors (e.g. quantized vectors). **Parameters** @@ -89,21 +85,20 @@ without creating intermediate heap arrays. | --- | --- | | `data` | The byte vectors (packed bits for binary quantization) | | `bytesPerVector` | The number of bytes in each vector | -| `resources` | The CuVS resources for device matrix creation | **Returns** -an instance of CuVSMatrix with BYTE data type +a host-memory CuVSMatrix with BYTE data type -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:92`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:80`_ ### createByteMatrixFromArray ```java -static CuVSMatrix createByteMatrixFromArray( byte[][] data, int bytesPerVector, CuVSResources resources) +static CuVSMatrix createByteMatrixFromArray(byte[][] data, int bytesPerVector) throws IOException ``` -A method to build a CuVSMatrix from a 2D byte array (for binary quantized vectors). +Builds a host-memory CuVSMatrix from a 2D byte array (e.g. quantized vectors). **Parameters** @@ -111,13 +106,33 @@ A method to build a CuVSMatrix from a 2D byte array (for binary quantized vector | --- | --- | | `data` | The 2D byte array (packed bits for binary quantization) | | `bytesPerVector` | The number of bytes in each vector | -| `resources` | The CuVS resources for device matrix creation | **Returns** -an instance of CuVSMatrix with BYTE data type +a host-memory CuVSMatrix with BYTE data type + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:98`_ + +### closeIndexWithDatasetFallback + +```java +static void closeIndexWithDatasetFallback(AutoCloseable index, AutoCloseable dataset) throws Exception +``` + +Closes an index that owns `dataset`. If index cleanup fails before releasing the +dataset, a direct dataset close is attempted and attached to the index failure when needed. + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:114`_ + +### ownDataset + +```java +static OwnedIndex ownDataset(AutoCloseable dataset) +``` + +Starts an ownership scope for a dataset that may later be transferred to an index. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:119`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:131`_ ### nanosToMillis @@ -137,7 +152,7 @@ A utility method to convert nanoseconds to milliseconds. milliseconds -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:141`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:196`_ ### cuVSResourcesOrNull @@ -151,7 +166,7 @@ Creates an instance of CuVSResources. an instance of CuVSResources -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:150`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:205`_ ### handleThrowableWithIgnore @@ -174,7 +189,7 @@ A utility method that conditionally ignores certain throwable objects | --- | --- | | `IOException` | | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:178`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:233`_ ### createListFromMergedVectors @@ -200,7 +215,7 @@ a list of float arrays | --- | --- | | `IOException` | I/O Exception | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:192`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:247`_ ### info @@ -218,6 +233,6 @@ Utility to print info/debug messages via InfoStream. | `component` | the name of the index writer | | `msg` | the log message to push via the InfoStream | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:210`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:265`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:26`_ diff --git a/fern/pages/other/multidimensional_arrays.md b/fern/pages/other/multidimensional_arrays.md index c5f7ec7d5e..072afa1087 100644 --- a/fern/pages/other/multidimensional_arrays.md +++ b/fern/pages/other/multidimensional_arrays.md @@ -165,16 +165,16 @@ import com.nvidia.cuvs.CuVSDeviceMatrix; long nRows = 100_000; long nFeatures = 128; -try (CuVSResources resources = CuVSResources.create()) { - CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, nRows, nFeatures, CuVSMatrix.DataType.FLOAT); - +try (CuVSResources resources = CuVSResources.create(); + CuVSMatrix.Builder builder = + CuVSMatrix.deviceBuilder( + resources, nRows, nFeatures, CuVSMatrix.DataType.FLOAT)) { for (long row = 0; row < nRows; row++) { builder.addVector(loadVector(row)); } - try (CuVSMatrix dataset = builder.build()) { + // A successful build transfers matrix ownership from the builder. + try (CuVSDeviceMatrix dataset = builder.build()) { // Pass dataset to NVIDIA cuVS Java APIs. } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java index 73fc0519a8..e0065cd91b 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java @@ -15,10 +15,8 @@ public interface CuVSDeviceMatrix extends CuVSMatrix { * responsible to call {@link CuVSMatrix#close()} to free its resources when done. */ default CuVSHostMatrix toHost() { - CuVSHostMatrix hostMatrix; - try (var builder = CuVSMatrix.hostBuilder(size(), columns(), dataType())) { - hostMatrix = builder.build(); - } + CuVSHostMatrix hostMatrix = + MatrixBuilderLifecycle.buildAndClose(CuVSMatrix.hostBuilder(size(), columns(), dataType())); try { toHost(hostMatrix); return hostMatrix; diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java index 3607fe0130..e8014d73e4 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java @@ -11,10 +11,9 @@ public interface CuVSHostMatrix extends CuVSMatrix { int get(int row, int col); default CuVSDeviceMatrix toDevice(CuVSResources resources) { - CuVSDeviceMatrix deviceMatrix; - try (var builder = CuVSMatrix.deviceBuilder(resources, size(), columns(), dataType())) { - deviceMatrix = builder.build(); - } + CuVSDeviceMatrix deviceMatrix = + MatrixBuilderLifecycle.buildAndClose( + CuVSMatrix.deviceBuilder(resources, size(), columns(), dataType())); try { toDevice(deviceMatrix, resources); return deviceMatrix; diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/MatrixBuilderLifecycle.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/MatrixBuilderLifecycle.java new file mode 100644 index 0000000000..c858535039 --- /dev/null +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/MatrixBuilderLifecycle.java @@ -0,0 +1,47 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +import java.util.Objects; + +/** Ownership-aware finalization shared by the public matrix conversion defaults. */ +final class MatrixBuilderLifecycle { + + private MatrixBuilderLifecycle() {} + + static T buildAndClose(CuVSMatrix.Builder builder) { + T matrix; + try { + matrix = Objects.requireNonNull(builder.build(), "Matrix builder must not return null"); + } catch (RuntimeException | Error operationFailure) { + closeAndSuppress(builder, operationFailure); + throw operationFailure; + } + + try { + builder.close(); + return matrix; + } catch (RuntimeException | Error builderCloseFailure) { + closeAndSuppress(matrix, builderCloseFailure); + throw builderCloseFailure; + } + } + + private static void closeAndSuppress(AutoCloseable resource, Throwable failure) { + try { + resource.close(); + } catch (RuntimeException | Error closeFailure) { + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } + } catch (Exception closeFailure) { + // CuVSMatrix and its Builder narrow close() to unchecked failures. Keep this guard so an + // unusual AutoCloseable implementation cannot replace the original failure. + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } + } + } +} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index c240e35d3e..eca0c96336 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -30,6 +30,8 @@ import java.util.List; import java.util.Locale; import java.util.Objects; +import java.util.function.Function; +import java.util.function.Supplier; import java.util.jar.JarFile; import java.util.jar.Manifest; import java.util.logging.Level; @@ -544,10 +546,15 @@ public CuVSMatrix.Builder newHostMatrixBuilder( @Override public CuVSMatrix.Builder newDeviceMatrixBuilder( CuVSResources resources, long size, long columns, CuVSMatrix.DataType dataType) { - var rowBytes = columns * dataType.bytes(); - return rowBytes > PinnedMemoryBuffer.CHUNK_BYTES - ? new DirectDeviceMatrixBuilder(resources, size, columns, dataType) - : new BufferedDeviceMatrixBuilder(resources, size, columns, dataType); + var rowBytes = deviceMatrixRowBytes(size, columns, dataType); + var bufferRowCount = Math.min(PinnedMemoryBuffer.CHUNK_BYTES / rowBytes, size); + return createDeviceMatrixBuilder( + () -> Util.getStream(resources), + stream -> + rowBytes > PinnedMemoryBuffer.CHUNK_BYTES + ? new DirectDeviceMatrixBuilder(resources, size, columns, dataType, stream) + : new BufferedDeviceMatrixBuilder( + resources, size, columns, dataType, stream, bufferRowCount)); } @Override @@ -558,11 +565,40 @@ public CuVSMatrix.Builder newDeviceMatrixBuilder( int rowStride, int columnStride, CuVSMatrix.DataType dataType) { - var rowBytes = columns * dataType.bytes(); - return rowBytes > PinnedMemoryBuffer.CHUNK_BYTES - ? new DirectDeviceMatrixBuilder(resources, size, columns, rowStride, columnStride, dataType) - : new BufferedDeviceMatrixBuilder( - resources, size, columns, rowStride, columnStride, dataType); + var rowBytes = deviceMatrixRowBytes(size, columns, dataType); + var bufferRowCount = Math.min(PinnedMemoryBuffer.CHUNK_BYTES / rowBytes, size); + return createDeviceMatrixBuilder( + () -> Util.getStream(resources), + stream -> + rowBytes > PinnedMemoryBuffer.CHUNK_BYTES + ? new DirectDeviceMatrixBuilder( + resources, size, columns, rowStride, columnStride, dataType, stream) + : new BufferedDeviceMatrixBuilder( + resources, + size, + columns, + rowStride, + columnStride, + dataType, + stream, + bufferRowCount)); + } + + private static long deviceMatrixRowBytes(long size, long columns, CuVSMatrix.DataType dataType) { + if (size < 0) { + throw new IllegalArgumentException("size must be non-negative: " + size); + } + if (columns <= 0) { + throw new IllegalArgumentException("columns must be positive: " + columns); + } + return Math.multiplyExact(columns, dataType.bytes()); + } + + /** Acquires the stream before invoking a constructor that allocates device memory. */ + static T createDeviceMatrixBuilder( + Supplier streamSupplier, Function builderFactory) { + MemorySegment stream = streamSupplier.get(); + return builderFactory.apply(stream); } @Override @@ -730,12 +766,17 @@ private static final class BufferedDeviceMatrixBuilder extends MatrixBuilder MatrixBuilderLifecycle.buildAndClose(builder)); + + assertSame(builderFailure, thrown); + assertEquals(1, builder.closeCalls.get()); + assertEquals(1, matrixCloses.get()); + } + + @Test + public void testMatrixCloseFailureIsSuppressedOnBuilderCloseFailure() { + IllegalStateException builderFailure = new IllegalStateException("builder close"); + IllegalArgumentException matrixFailure = new IllegalArgumentException("matrix close"); + AtomicInteger matrixCloses = new AtomicInteger(); + CuVSMatrix matrix = fakeMatrix(matrixCloses, matrixFailure); + FakeBuilder builder = new FakeBuilder(matrix, null, builderFailure); + + IllegalStateException thrown = + assertThrows( + IllegalStateException.class, () -> MatrixBuilderLifecycle.buildAndClose(builder)); + + assertSame(builderFailure, thrown); + assertArrayEquals(new Throwable[] {matrixFailure}, thrown.getSuppressed()); + assertEquals(1, matrixCloses.get()); + } + + @Test + public void testBuildFailureRemainsPrimaryWhenBuilderCloseAlsoFails() { + IllegalArgumentException buildFailure = new IllegalArgumentException("build"); + IllegalStateException closeFailure = new IllegalStateException("builder close"); + FakeBuilder builder = new FakeBuilder(null, buildFailure, closeFailure); + + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, () -> MatrixBuilderLifecycle.buildAndClose(builder)); + + assertSame(buildFailure, thrown); + assertArrayEquals(new Throwable[] {closeFailure}, thrown.getSuppressed()); + assertEquals(1, builder.closeCalls.get()); + } + + private static CuVSMatrix fakeMatrix(AtomicInteger closeCalls, RuntimeException closeFailure) { + return (CuVSMatrix) + Proxy.newProxyInstance( + CuVSMatrix.class.getClassLoader(), + new Class[] {CuVSMatrix.class}, + (proxy, method, args) -> { + if (method.getName().equals("close") && method.getParameterCount() == 0) { + closeCalls.incrementAndGet(); + if (closeFailure != null) { + throw closeFailure; + } + return null; + } + throw new AssertionError("Unexpected matrix call: " + method); + }); + } + + private static final class FakeBuilder implements CuVSMatrix.Builder { + private final CuVSMatrix matrix; + private final RuntimeException buildFailure; + private final RuntimeException closeFailure; + private final AtomicInteger closeCalls = new AtomicInteger(); + + private FakeBuilder( + CuVSMatrix matrix, RuntimeException buildFailure, RuntimeException closeFailure) { + this.matrix = matrix; + this.buildFailure = buildFailure; + this.closeFailure = closeFailure; + } + + @Override + public void addVector(float[] vector) {} + + @Override + public void addVector(byte[] vector) {} + + @Override + public void addVector(int[] vector) {} + + @Override + public void addVector(short[] vector) {} + + @Override + public CuVSMatrix build() { + if (buildFailure != null) { + throw buildFailure; + } + return matrix; + } + + @Override + public void close() { + closeCalls.incrementAndGet(); + if (closeFailure != null) { + throw closeFailure; + } + } + } +} diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java index 076db40553..b9ed704d70 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java @@ -6,11 +6,18 @@ import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; import com.nvidia.cuvs.CuVSTestCase; +import com.nvidia.cuvs.LibraryException; +import java.lang.foreign.MemorySegment; import java.lang.invoke.MethodHandles; import java.lang.invoke.MethodType; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import java.util.function.Supplier; import org.junit.Before; import org.junit.Test; @@ -118,6 +125,57 @@ public void testMaxVersionOverride() { } } + @Test + public void testDeviceMatrixBuilderDoesNotAllocateWhenStreamAcquisitionFails() { + LibraryException streamFailure = new LibraryException("stream acquisition"); + AtomicInteger streamCalls = new AtomicInteger(); + AtomicInteger builderCalls = new AtomicInteger(); + + LibraryException thrown = + assertThrows( + LibraryException.class, + () -> + createDeviceMatrixBuilder( + () -> { + streamCalls.incrementAndGet(); + throw streamFailure; + }, + ignored -> { + builderCalls.incrementAndGet(); + return new Object(); + })); + + assertSame(streamFailure, thrown); + assertEquals(1, streamCalls.get()); + assertEquals(0, builderCalls.get()); + } + + @Test + public void testDeviceMatrixBuilderPassesAcquiredStreamToConstructor() throws Throwable { + MemorySegment stream = MemorySegment.ofArray(new byte[1]); + Object builder = new Object(); + AtomicInteger streamCalls = new AtomicInteger(); + AtomicInteger builderCalls = new AtomicInteger(); + AtomicReference observedStream = new AtomicReference<>(); + + Object created = + createDeviceMatrixBuilder( + () -> { + streamCalls.incrementAndGet(); + return stream; + }, + actualStream -> { + builderCalls.incrementAndGet(); + observedStream.set(actualStream); + return builder; + }); + + assertSame(builder, created); + assertSame(stream, observedStream.get()); + assertEquals(1, streamCalls.get()); + assertEquals(1, builderCalls.get()); + } + static void checkCuVSVersionMatching(String mavenVersionString, int major, int minor, int patch) throws ProviderInitializationException { try { @@ -136,4 +194,17 @@ static void checkCuVSVersionMatching(String mavenVersionString, int major, int m throw new AssertionError(e); } } + + static Object createDeviceMatrixBuilder( + Supplier streamSupplier, Function builderFactory) + throws Throwable { + var cls = Class.forName("com.nvidia.cuvs.spi.JDKProvider"); + var method = + MethodHandles.lookup() + .findStatic( + cls, + "createDeviceMatrixBuilder", + MethodType.methodType(Object.class, Supplier.class, Function.class)); + return method.invoke(streamSupplier, builderFactory); + } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 31c2213492..a4c96b9df6 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -299,15 +299,16 @@ static CuVSMatrix remapSubsetGraph(CuVSMatrix cagraGraph, int[] selectedNodes) throw new IOException("The subset graph must have a positive degree"); } int[] remappedRow = new int[degree]; - try (CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(selectedNodes.length, degree, CuVSMatrix.DataType.INT)) { - for (int rowOrdinal = 0; rowOrdinal < selectedNodes.length; rowOrdinal++) { - RowView row = cagraGraph.getRow(rowOrdinal); - remapSubsetAdjacencyRow(rowOrdinal, row, degree, selectedNodes, remappedRow); - builder.addVector(remappedRow); - } - return builder.build(); - } + return MatrixBuilderLifecycle.build( + CuVSMatrix.hostBuilder(selectedNodes.length, degree, CuVSMatrix.DataType.INT), + builder -> { + for (int rowOrdinal = 0; rowOrdinal < selectedNodes.length; rowOrdinal++) { + RowView row = cagraGraph.getRow(rowOrdinal); + remapSubsetAdjacencyRow(rowOrdinal, row, degree, selectedNodes, remappedRow); + builder.addVector(remappedRow); + } + return builder.build(); + }); } static void remapSubsetAdjacencyRow( @@ -346,51 +347,54 @@ static void remapSubsetAdjacencyRow( private static CuVSMatrix createSubsetDataset( CuVSMatrix vectors, int[] selectedNodes, int columns, CuVSMatrix.DataType dataType) throws IOException { - try (CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(selectedNodes.length, columns, dataType)) { - if (dataType == CuVSMatrix.DataType.FLOAT) { - float[] rowBuffer = new float[columns]; - for (int node : selectedNodes) { - RowView row = vectors.getRow(node); - validateRowWidth(row, columns, node); - row.toArray(rowBuffer); - builder.addVector(rowBuffer); - } - } else { - byte[] rowBuffer = new byte[columns]; - for (int node : selectedNodes) { - RowView row = vectors.getRow(node); - validateRowWidth(row, columns, node); - row.toArray(rowBuffer); - builder.addVector(rowBuffer); - } - } - return builder.build(); - } + return MatrixBuilderLifecycle.build( + CuVSMatrix.hostBuilder(selectedNodes.length, columns, dataType), + builder -> { + if (dataType == CuVSMatrix.DataType.FLOAT) { + float[] rowBuffer = new float[columns]; + for (int node : selectedNodes) { + RowView row = vectors.getRow(node); + validateRowWidth(row, columns, node); + row.toArray(rowBuffer); + builder.addVector(rowBuffer); + } + } else { + byte[] rowBuffer = new byte[columns]; + for (int node : selectedNodes) { + RowView row = vectors.getRow(node); + validateRowWidth(row, columns, node); + row.toArray(rowBuffer); + builder.addVector(rowBuffer); + } + } + return builder.build(); + }); } private static CuVSMatrix createSubsetDataset( - List vectors, int[] selectedNodes, int columns, CuVSMatrix.DataType dataType) { - try (CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(selectedNodes.length, columns, dataType)) { - for (int node : selectedNodes) { - Object vector = vectors.get(node); - if (dataType == CuVSMatrix.DataType.FLOAT) { - if (!(vector instanceof float[] values) || values.length != columns) { - throw new IllegalArgumentException( - "Vector " + node + " must be a float[" + columns + "]"); - } - builder.addVector(values); - } else { - if (!(vector instanceof byte[] values) || values.length != columns) { - throw new IllegalArgumentException( - "Vector " + node + " must be a byte[" + columns + "]"); + List vectors, int[] selectedNodes, int columns, CuVSMatrix.DataType dataType) + throws IOException { + return MatrixBuilderLifecycle.build( + CuVSMatrix.hostBuilder(selectedNodes.length, columns, dataType), + builder -> { + for (int node : selectedNodes) { + Object vector = vectors.get(node); + if (dataType == CuVSMatrix.DataType.FLOAT) { + if (!(vector instanceof float[] values) || values.length != columns) { + throw new IllegalArgumentException( + "Vector " + node + " must be a float[" + columns + "]"); + } + builder.addVector(values); + } else { + if (!(vector instanceof byte[] values) || values.length != columns) { + throw new IllegalArgumentException( + "Vector " + node + " must be a byte[" + columns + "]"); + } + builder.addVector(values); + } } - builder.addVector(values); - } - } - return builder.build(); - } + return builder.build(); + }); } private static void validateRowWidth(RowView row, int columns, int rowIndex) throws IOException { diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 4460ee9187..94e3626065 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -20,6 +20,7 @@ import com.nvidia.cuvs.BruteForceIndexParams; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CuVSDeviceMatrix; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.LibraryException; import java.io.IOException; @@ -202,8 +203,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro if (indexType.isCagra()) { var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); try { - CuVSMatrix cagraDataset = - Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + CuVSDeviceMatrix cagraDataset = + createDeviceFloatMatrix(vectors, fieldInfo.getVectorDimension()); writeCagraIndex(cagraIndexOutputStream, cagraDataset); } catch (RecoverableCagraConstructionException recoverable) { ensureFallbackOutputUnchanged(cagraIndexOffset, cuvsIndex.getFilePointer(), recoverable); @@ -239,16 +240,79 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro } } + /** + * Creates GPU-search input directly on the device. Accelerated HNSW intentionally continues to + * use the shared host-backed matrix helpers because it does not retain a searchable device + * dataset. + */ + private CuVSDeviceMatrix createDeviceFloatMatrix(List vectors, int dimensions) + throws Throwable { + return prepareCagraDataset( + () -> + CuVSMatrix.deviceBuilder( + getCuVSResourcesInstance(), vectors.size(), dimensions, CuVSMatrix.DataType.FLOAT), + builder -> { + for (float[] vector : vectors) { + builder.addVector(vector); + } + }); + } + + @FunctionalInterface + interface CagraDatasetBuilderFactory { + CuVSMatrix.Builder create() throws Throwable; + } + + @FunctionalInterface + interface CagraDatasetPopulation { + void populate(CuVSMatrix.Builder builder) throws Throwable; + } + + /** + * Builds the device input and applies the pre-persistence fallback policy only after builder + * cleanup succeeds. A builder cleanup failure is fatal; after ownership has transferred, this + * method also closes the dataset before propagating that failure. A builder-factory failure is + * also fatal because no cleanup handle was returned. + */ + static CuVSDeviceMatrix prepareCagraDataset( + CagraDatasetBuilderFactory builderFactory, CagraDatasetPopulation population) + throws Throwable { + CuVSMatrix.Builder builder = null; + CuVSDeviceMatrix dataset = null; + Throwable operationFailure = null; + boolean factoryCompleted = false; + try { + builder = + Objects.requireNonNull(builderFactory.create(), "CAGRA dataset builder must not be null"); + factoryCompleted = true; + population.populate(builder); + dataset = builder.build(); + } catch (Throwable failure) { + operationFailure = failure; + } + + Throwable cleanupFailure = closeResource(builder, null); + if (cleanupFailure != null && dataset != null) { + cleanupFailure = closeResource(dataset, cleanupFailure); + } + + Throwable failure = combineOperationAndCleanupFailures(operationFailure, cleanupFailure); + if (failure != null) { + throw classifyCagraWriteFailure( + failure, false, cleanupFailure != null || factoryCompleted == false); + } + return dataset; + } + /** * Builds and writes the CAGRA index. * * @param os Instance of the OutputStream - * @param dataset The instance of CuVSMatrix holding the dataset + * @param dataset device-backed matrix holding the dataset * @throws Throwable */ - private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { + private void writeCagraIndex(OutputStream os, CuVSDeviceMatrix dataset) throws Throwable { CagraIndex index = null; - CuVSMatrix deviceVectors = null; AutoCloseable indexDataset = null; Throwable failure = null; CagraWriteContext writeContext = new CagraWriteContext(); @@ -260,17 +324,16 @@ private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwab .withDataset(dataset) .withIndexParams(params) .build(); - deviceVectors = dataset.toDevice(getCuVSResourcesInstance()); /* * cuVS rejects makePaddedDataset for a device matrix whose rows already sit at the required * stride, and asks for a view over that storage instead. */ - if (CagraIndex.isPaddedDataset(deviceVectors)) { - var indexDatasetView = index.makePaddedDatasetView(deviceVectors); + if (CagraIndex.isPaddedDataset(dataset)) { + var indexDatasetView = index.makePaddedDatasetView(dataset); indexDataset = indexDatasetView; index.updateDataset(indexDatasetView); } else { - var paddedDataset = index.makePaddedDataset(deviceVectors); + var paddedDataset = index.makePaddedDataset(dataset); indexDataset = paddedDataset; index.updateDataset(paddedDataset); } @@ -283,7 +346,7 @@ private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwab failure = t; } - Throwable cleanupFailure = closeCagraResources(index, dataset, indexDataset, deviceVectors); + Throwable cleanupFailure = closeCagraResources(index, dataset, indexDataset); failure = combineOperationAndCleanupFailures(failure, cleanupFailure); if (failure != null) { throw classifyCagraWriteFailure( @@ -327,10 +390,7 @@ private static Throwable closeResource(AutoCloseable resource, Throwable failure /** Closes CAGRA resources in dependency order and returns the first cleanup failure. */ static Throwable closeCagraResources( - AutoCloseable index, - AutoCloseable originalDataset, - AutoCloseable indexDataset, - AutoCloseable deviceVectors) { + AutoCloseable index, AutoCloseable originalDataset, AutoCloseable indexDataset) { Throwable failure = null; try { if (index == null) { @@ -341,8 +401,7 @@ static Throwable closeCagraResources( } catch (Throwable closeFailure) { failure = addFailure(failure, closeFailure); } - failure = closeResource(indexDataset, failure); - return closeResource(deviceVectors, failure); + return closeResource(indexDataset, failure); } static Throwable combineOperationAndCleanupFailures( diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MatrixBuilderLifecycle.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MatrixBuilderLifecycle.java new file mode 100644 index 0000000000..4712c78ac3 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MatrixBuilderLifecycle.java @@ -0,0 +1,64 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CuVSMatrix; +import java.io.IOException; +import java.util.Objects; + +/** Ownership-aware matrix construction shared by the Lucene ingestion paths. */ +final class MatrixBuilderLifecycle { + + private MatrixBuilderLifecycle() {} + + @FunctionalInterface + interface BuildOperation { + T build(CuVSMatrix.Builder builder) throws Throwable; + } + + static T build(CuVSMatrix.Builder builder, BuildOperation operation) + throws IOException { + T matrix = null; + Throwable operationFailure = null; + try { + matrix = + Objects.requireNonNull(operation.build(builder), "Matrix builder must not return null"); + } catch (Throwable failure) { + operationFailure = failure; + } + + Throwable cleanupFailure = closeResource(builder, null); + if (cleanupFailure != null && matrix != null) { + cleanupFailure = closeResource(matrix, cleanupFailure); + } + Throwable failure = addFailure(operationFailure, cleanupFailure); + if (failure != null) { + throw Utils.handleThrowable(failure); + } + return matrix; + } + + private static Throwable closeResource(AutoCloseable resource, Throwable failure) { + try { + resource.close(); + } catch (Throwable closeFailure) { + return addFailure(failure, closeFailure); + } + return failure; + } + + private static Throwable addFailure(Throwable primary, Throwable secondary) { + if (secondary == null) { + return primary; + } + if (primary == null) { + return secondary; + } + if (primary != secondary) { + primary.addSuppressed(secondary); + } + return primary; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java index a47f65658c..cdc1459e71 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -60,14 +60,15 @@ static RuntimeException handleThrowable(Throwable t) throws IOException { * @param dimensions The number of float elements in each vector * @return a host-memory CuVSMatrix */ - static CuVSMatrix createFloatMatrix(List data, int dimensions) { - try (CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT)) { - for (float[] vector : data) { - builder.addVector(vector); - } - return builder.build(); - } + static CuVSMatrix createFloatMatrix(List data, int dimensions) throws IOException { + return MatrixBuilderLifecycle.build( + CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT), + builder -> { + for (float[] vector : data) { + builder.addVector(vector); + } + return builder.build(); + }); } /** @@ -77,14 +78,15 @@ static CuVSMatrix createFloatMatrix(List data, int dimensions) { * @param bytesPerVector The number of bytes in each vector * @return a host-memory CuVSMatrix with BYTE data type */ - static CuVSMatrix createByteMatrix(List data, int bytesPerVector) { - try (CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE)) { - for (byte[] vector : data) { - builder.addVector(vector); - } - return builder.build(); - } + static CuVSMatrix createByteMatrix(List data, int bytesPerVector) throws IOException { + return MatrixBuilderLifecycle.build( + CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE), + builder -> { + for (byte[] vector : data) { + builder.addVector(vector); + } + return builder.build(); + }); } /** @@ -94,14 +96,16 @@ static CuVSMatrix createByteMatrix(List data, int bytesPerVector) { * @param bytesPerVector The number of bytes in each vector * @return a host-memory CuVSMatrix with BYTE data type */ - static CuVSMatrix createByteMatrixFromArray(byte[][] data, int bytesPerVector) { - try (CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE)) { - for (byte[] vector : data) { - builder.addVector(vector); - } - return builder.build(); - } + static CuVSMatrix createByteMatrixFromArray(byte[][] data, int bytesPerVector) + throws IOException { + return MatrixBuilderLifecycle.build( + CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE), + builder -> { + for (byte[] vector : data) { + builder.addVector(vector); + } + return builder.build(); + }); } /** diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java index eef332bd85..9cc5ec87ed 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java @@ -7,6 +7,7 @@ import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; import static com.nvidia.cuvs.lucene.TestUtils.generateRandomVector; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; import java.io.IOException; import java.util.ArrayList; @@ -19,14 +20,22 @@ import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.HnswGraphProvider; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; import org.apache.lucene.document.Document; import org.apache.lucene.document.Field; import org.apache.lucene.document.KnnFloatVectorField; import org.apache.lucene.document.StringField; +import org.apache.lucene.index.CodecReader; import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.index.NoMergePolicy; import org.apache.lucene.index.Term; +import org.apache.lucene.index.TieredMergePolicy; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.KnnFloatVectorQuery; @@ -41,6 +50,7 @@ import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.hnsw.HnswGraph; import org.junit.BeforeClass; import org.junit.Test; @@ -321,13 +331,14 @@ public void testVectorSearchWithPartialDeletionAndReindexing() throws IOExceptio } @Test - public void testForceMergeCountsOnlyLiveSparseVectors() throws IOException { + public void testForceMergeCountsOnlyLiveSparseVectors() throws Exception { final String vectorField = "vector"; final int dimensions = 129; Map expected = new LinkedHashMap<>(); try (Directory directory = newDirectory()) { - try (IndexWriter writer = new IndexWriter(directory, createWriterConfig())) { + try (IndexWriter writer = + new IndexWriter(directory, createWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))) { for (int segment = 0; segment < 3; segment++) { for (int row = 0; row < 5; row++) { String id = segment + "-" + row; @@ -349,6 +360,19 @@ public void testForceMergeCountsOnlyLiveSparseVectors() throws IOException { writer.deleteDocuments(new Term("id", segment + "-1")); } writer.commit(); + + try (DirectoryReader sourceReader = DirectoryReader.open(writer)) { + assertEquals("the test requires three source segments", 3, sourceReader.leaves().size()); + for (var context : sourceReader.leaves()) { + LeafReader sourceLeaf = context.reader(); + assertTrue("each source segment must carry a deletion", sourceLeaf.hasDeletions()); + assertEquals(5, sourceLeaf.maxDoc()); + assertEquals(4, sourceLeaf.numDocs()); + assertEquals(4, sourceLeaf.getFloatVectorValues(vectorField).size()); + } + } + + writer.getConfig().setMergePolicy(new TieredMergePolicy()); writer.forceMerge(1); } @@ -366,6 +390,116 @@ public void testForceMergeCountsOnlyLiveSparseVectors() throws IOException { assertArrayEquals(expected.get(id), values.vectorValue(ordinal), 0.0f); } assertEquals(expected.keySet(), seen); + + HnswGraph graph = graphOf(leaf, vectorField); + assertEquals(values.size(), graph.size()); + assertEquals(values.size(), graph.getNodesOnLevel(0).size()); + assertAllGraphOrdinalsInBounds(graph, values.size()); + } + } + } + + @Test + public void testForceMergeWithExactlyOneLiveVector() throws Exception { + assertTrivialLiveVectorMerge(1); + } + + @Test + public void testForceMergeWithZeroLiveVectors() throws Exception { + assertTrivialLiveVectorMerge(0); + } + + private void assertTrivialLiveVectorMerge(int liveVectors) throws Exception { + final String vectorField = "vector"; + final int dimensions = 129; + + try (Directory directory = newDirectory()) { + try (IndexWriter writer = + new IndexWriter(directory, createWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))) { + for (int id = 0; id < 3; id++) { + Document vectorDocument = new Document(); + vectorDocument.add(new StringField("id", "vector-" + id, Field.Store.YES)); + vectorDocument.add( + new KnnFloatVectorField( + vectorField, + deterministicVector(id, dimensions), + VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(vectorDocument); + + Document sparseDocument = new Document(); + sparseDocument.add(new StringField("id", "sparse-" + id, Field.Store.YES)); + writer.addDocument(sparseDocument); + writer.commit(); + } + for (int id = liveVectors; id < 3; id++) { + writer.deleteDocuments(new Term("id", "vector-" + id)); + } + writer.commit(); + + try (DirectoryReader sourceReader = DirectoryReader.open(writer)) { + assertEquals("the test requires three source segments", 3, sourceReader.leaves().size()); + } + + writer.getConfig().setMergePolicy(new TieredMergePolicy()); + writer.forceMerge(1); + } + + TestUtil.checkIndex(directory); + try (DirectoryReader reader = DirectoryReader.open(directory)) { + LeafReader leaf = getOnlyLeafReader(reader); + FloatVectorValues values = leaf.getFloatVectorValues(vectorField); + assertNotNull(values); + assertEquals(liveVectors, values.size()); + + HnswGraph graph = graphOf(leaf, vectorField); + assertEquals(liveVectors, graph.size()); + assertEquals(liveVectors == 0 ? 0 : 1, graph.numLevels()); + assertEquals(liveVectors, graph.getNodesOnLevel(0).size()); + assertAllGraphOrdinalsInBounds(graph, liveVectors); + if (liveVectors == 1) { + graph.seek(0, 0); + assertEquals( + "a single-node graph must have no edges", NO_MORE_DOCS, graph.nextNeighbor()); + } + + ((CodecReader) leaf).getVectorReader().checkIntegrity(); + IndexSearcher searcher = new IndexSearcher(reader); + TopDocs results = + searcher.search( + new KnnFloatVectorQuery(vectorField, deterministicVector(0, dimensions), 1), 1); + assertEquals(liveVectors, results.totalHits.value()); + assertEquals(liveVectors, results.scoreDocs.length); + if (liveVectors == 1) { + assertEquals( + "vector-0", searcher.storedFields().document(results.scoreDocs[0].doc).get("id")); + } + } + } + } + + private static HnswGraph graphOf(LeafReader leaf, String field) throws Exception { + KnnVectorsReader reader = ((CodecReader) leaf).getVectorReader(); + if (reader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + reader = fieldsReader.getFieldReader(field); + } + return ((HnswGraphProvider) reader).getGraph(field); + } + + private static void assertAllGraphOrdinalsInBounds(HnswGraph graph, int vectorCount) + throws Exception { + for (int level = 0; level < graph.numLevels(); level++) { + HnswGraph.NodesIterator nodes = graph.getNodesOnLevel(level); + while (nodes.hasNext()) { + int node = nodes.nextInt(); + assertTrue("graph node is outside the vector domain", node >= 0 && node < vectorCount); + graph.seek(level, node); + for (int neighbor = graph.nextNeighbor(); + neighbor != NO_MORE_DOCS; + neighbor = graph.nextNeighbor()) { + assertTrue( + "graph neighbor is outside the vector domain", + neighbor >= 0 && neighbor < vectorCount); + } } } } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java index 16a5c8b84d..084c58a253 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java @@ -104,14 +104,17 @@ public void testThreeLevelGraphCanBeReopenedAndSearched() throws Exception { assertTrue(nodesByLevel.get(1).containsAll(nodesByLevel.get(2))); assertUpperNeighborsStayOnTheirLevel(graph, nodesByLevel); + int queryNode = graph.entryNode(); + assertTrue(nodesByLevel.get(2).contains(queryNode)); IndexSearcher searcher = new IndexSearcher(reader); - var hits = searcher.search(new KnnFloatVectorQuery(FIELD, vectors[0], 10), 10); + var hits = searcher.search(new KnnFloatVectorQuery(FIELD, vectors[queryNode], 10), 10); assertEquals(10, hits.scoreDocs.length); - boolean foundExactVector = false; + String queryNodeId = Integer.toString(queryNode); + boolean foundQueryNode = false; for (var hit : hits.scoreDocs) { - foundExactVector |= "0".equals(searcher.storedFields().document(hit.doc).get("id")); + foundQueryNode |= queryNodeId.equals(searcher.storedFields().document(hit.doc).get("id")); } - assertTrue("the indexed vector must be returned for its own query", foundExactVector); + assertTrue("the entry-node vector must be returned for its own query", foundQueryNode); } } } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java index 89d58aa682..5effa973b1 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java @@ -11,10 +11,15 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; import org.apache.lucene.document.Document; import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DirectoryReader; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.util.LuceneTestCase; @@ -38,14 +43,14 @@ public class TestCagraIndexAtAlignedDimensions extends LuceneTestCase { @Test public void testCagraIsBuiltAtAnAlignedDimension() throws IOException { - // 128 floats is 512 bytes, an exact multiple of the 16 byte CAGRA row alignment. - assertCagraIsBuilt(128); + // Deep1B's 96 floats occupy 384 bytes, an exact multiple of CAGRA's 16-byte row alignment. + assertCagraIsBuilt(96); } @Test public void testCagraIsBuiltAtAnUnalignedDimension() throws IOException { - // 127 floats is not, so the writer has to fall back to an owning padded copy. - assertCagraIsBuilt(127); + // 95 floats is not aligned, so the writer has to create an owning padded copy. + assertCagraIsBuilt(95); } /** Indexes a segment of the given dimension and fails if the CAGRA build did not survive it. */ @@ -54,6 +59,11 @@ private void assertCagraIsBuilt(int dimension) throws IOException { RecordingInfoStream infoStream = new RecordingInfoStream(); try (Directory directory = newDirectory()) { + float[] queryVector = new float[dimension]; + for (int d = 0; d < dimension; d++) { + queryVector[d] = random().nextFloat(); + } + IndexWriterConfig config = new IndexWriterConfig() .setCodec( @@ -64,9 +74,11 @@ private void assertCagraIsBuilt(int dimension) throws IOException { try (IndexWriter writer = new IndexWriter(directory, config)) { for (int i = 0; i < 64; i++) { - float[] vector = new float[dimension]; - for (int d = 0; d < dimension; d++) { - vector[d] = random().nextFloat(); + float[] vector = i == 0 ? queryVector : new float[dimension]; + if (i != 0) { + for (int d = 0; d < dimension; d++) { + vector[d] = random().nextFloat(); + } } Document doc = new Document(); doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.EUCLIDEAN)); @@ -74,6 +86,24 @@ private void assertCagraIsBuilt(int dimension) throws IOException { } writer.commit(); } + + try (DirectoryReader reader = DirectoryReader.open(directory)) { + LeafReader leaf = getOnlyLeafReader(reader); + CuVS2510GPUVectorsReader gpuReader = gpuReader(leaf, "vector"); + CuVS2510GPUVectorsReader.FieldEntry fieldEntry = gpuReader.getFieldEntry("vector"); + assertNotNull(fieldEntry); + assertTrue("Expected a serialized CAGRA payload", fieldEntry.cagraIndexLength() > 0); + assertEquals("Did not expect a brute-force payload", 0, fieldEntry.bruteForceIndexLength()); + assertNotNull( + "Expected the CAGRA index to be loaded", gpuReader.getCagraIndexForField("vector")); + + var searcher = newSearcher(reader); + var query = new GPUKnnFloatVectorQuery("vector", queryVector, 1, null, 1, 1); + var topDocs = searcher.search(query, 1); + assertEquals(1, topDocs.scoreDocs.length); + assertTrue(topDocs.scoreDocs[0].doc >= 0); + assertTrue(topDocs.scoreDocs[0].doc < reader.maxDoc()); + } } assertTrue( @@ -84,6 +114,14 @@ private void assertCagraIsBuilt(int dimension) throws IOException { infoStream.cagraBuildFailures().isEmpty()); } + private static CuVS2510GPUVectorsReader gpuReader(LeafReader leaf, String field) { + KnnVectorsReader reader = ((CodecReader) leaf).getVectorReader(); + if (reader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + reader = fieldsReader.getFieldReader(field); + } + return (CuVS2510GPUVectorsReader) reader; + } + /** An InfoStream that keeps the messages, so that a test can tell which index type was built. */ private static class RecordingInfoStream extends InfoStream { diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsWriterFailureHandling.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsWriterFailureHandling.java index 5062a56799..605ca8e64f 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsWriterFailureHandling.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsWriterFailureHandling.java @@ -4,8 +4,11 @@ */ package com.nvidia.cuvs.lucene; +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.LibraryException; import java.io.IOException; +import java.lang.reflect.Proxy; import java.util.ArrayList; import java.util.List; import java.util.concurrent.atomic.AtomicInteger; @@ -14,6 +17,140 @@ public class TestCuVS2510GPUVectorsWriterFailureHandling extends LuceneTestCase { + @Test + public void testDevicePreparationNativeFailureIsRecoverable() { + LibraryException nativeFailure = new LibraryException("device upload"); + FakeDeviceBuilder builder = new FakeDeviceBuilder(null, null); + + Throwable thrown = + assertThrows( + Throwable.class, + () -> + CuVS2510GPUVectorsWriter.prepareCagraDataset( + () -> builder, + ignored -> { + throw nativeFailure; + })); + + assertTrue(thrown instanceof CuVS2510GPUVectorsWriter.RecoverableCagraConstructionException); + assertSame(nativeFailure, thrown.getCause()); + assertEquals(0, nativeFailure.getSuppressed().length); + assertEquals(1, builder.closeCalls.get()); + } + + @Test + public void testDevicePreparationOperationAndCleanupFailureIsFatal() { + LibraryException nativeFailure = new LibraryException("device upload"); + IllegalStateException cleanupFailure = new IllegalStateException("builder close"); + FakeDeviceBuilder builder = new FakeDeviceBuilder(null, cleanupFailure); + + Throwable thrown = + assertThrows( + Throwable.class, + () -> + CuVS2510GPUVectorsWriter.prepareCagraDataset( + () -> builder, + ignored -> { + throw nativeFailure; + })); + + assertSame(nativeFailure, thrown); + assertArrayEquals(new Throwable[] {cleanupFailure}, thrown.getSuppressed()); + assertEquals(1, builder.closeCalls.get()); + } + + @Test + public void testDevicePreparationNativeFactoryFailureIsFatal() { + LibraryException factoryFailure = new LibraryException("builder factory"); + FakeDeviceBuilder builder = new FakeDeviceBuilder(null, null); + + Throwable thrown = + assertThrows( + Throwable.class, + () -> + CuVS2510GPUVectorsWriter.prepareCagraDataset( + () -> { + throw factoryFailure; + }, + ignored -> fail("population must not run"))); + + assertSame(factoryFailure, thrown); + assertFalse(thrown instanceof CuVS2510GPUVectorsWriter.RecoverableCagraConstructionException); + assertEquals(0, builder.closeCalls.get()); + } + + @Test + public void testDevicePreparationTransfersOwnershipAfterBuilderCleanup() throws Throwable { + AtomicInteger datasetCloses = new AtomicInteger(); + CuVSDeviceMatrix dataset = fakeDeviceMatrix(datasetCloses, null); + FakeDeviceBuilder builder = new FakeDeviceBuilder(dataset, null); + + CuVSDeviceMatrix prepared = + CuVS2510GPUVectorsWriter.prepareCagraDataset(() -> builder, ignored -> {}); + + assertSame(dataset, prepared); + assertEquals(1, builder.closeCalls.get()); + assertEquals(0, datasetCloses.get()); + } + + @Test + public void testDevicePreparationBuilderCleanupFailureClosesTransferredDataset() { + LibraryException cleanupFailure = new LibraryException("builder close"); + AtomicInteger datasetCloses = new AtomicInteger(); + CuVSDeviceMatrix dataset = fakeDeviceMatrix(datasetCloses, null); + FakeDeviceBuilder builder = new FakeDeviceBuilder(dataset, cleanupFailure); + + Throwable thrown = + assertThrows( + Throwable.class, + () -> CuVS2510GPUVectorsWriter.prepareCagraDataset(() -> builder, ignored -> {})); + + assertSame(cleanupFailure, thrown); + assertEquals(1, builder.closeCalls.get()); + assertEquals(1, datasetCloses.get()); + } + + @Test + public void testDevicePreparationDatasetCleanupFailureIsSuppressed() { + LibraryException builderFailure = new LibraryException("builder close"); + IllegalStateException datasetFailure = new IllegalStateException("dataset close"); + AtomicInteger datasetCloses = new AtomicInteger(); + CuVSDeviceMatrix dataset = fakeDeviceMatrix(datasetCloses, datasetFailure); + FakeDeviceBuilder builder = new FakeDeviceBuilder(dataset, builderFailure); + + Throwable thrown = + assertThrows( + Throwable.class, + () -> CuVS2510GPUVectorsWriter.prepareCagraDataset(() -> builder, ignored -> {})); + + assertSame(builderFailure, thrown); + assertArrayEquals(new Throwable[] {datasetFailure}, thrown.getSuppressed()); + assertEquals(1, builder.closeCalls.get()); + assertEquals(1, datasetCloses.get()); + } + + @Test + public void testDevicePreparationPreservesNonNativeFailures() { + IOException checkedFailure = new IOException("checked operation"); + AssertionError error = new AssertionError("operation error"); + + for (Throwable expected : List.of(checkedFailure, error)) { + FakeDeviceBuilder builder = new FakeDeviceBuilder(null, null); + Throwable thrown = + assertThrows( + Throwable.class, + () -> + CuVS2510GPUVectorsWriter.prepareCagraDataset( + () -> builder, + ignored -> { + throw expected; + })); + + assertSame(expected, thrown); + assertEquals(1, builder.closeCalls.get()); + } + } + @Test public void testNativeConstructionFailureIsRecoverableBeforePersistence() { LibraryException nativeFailure = new LibraryException("native construction"); @@ -106,11 +243,10 @@ public void testCagraResourcesCloseInDependencyOrder() { CuVS2510GPUVectorsWriter.closeCagraResources( () -> closeOrder.add("index"), directDatasetCloses::incrementAndGet, - () -> closeOrder.add("padded-dataset"), - () -> closeOrder.add("device-vectors")); + () -> closeOrder.add("padded-dataset")); assertNull(failure); - assertEquals(List.of("index", "padded-dataset", "device-vectors"), closeOrder); + assertEquals(List.of("index", "padded-dataset"), closeOrder); assertEquals(0, directDatasetCloses.get()); } @@ -121,23 +257,20 @@ public void testBodyAndCleanupFailuresPreserveOrderAndSuppression() { RuntimeException indexFailure = new RuntimeException("index close"); RuntimeException datasetFailure = new RuntimeException("dataset close"); RuntimeException paddedFailure = new RuntimeException("padded close"); - RuntimeException deviceFailure = new RuntimeException("device close"); Throwable cleanupFailure = CuVS2510GPUVectorsWriter.closeCagraResources( failingCloseable("index", closeOrder, indexFailure), failingCloseable("dataset", closeOrder, datasetFailure), - failingCloseable("padded-dataset", closeOrder, paddedFailure), - failingCloseable("device-vectors", closeOrder, deviceFailure)); + failingCloseable("padded-dataset", closeOrder, paddedFailure)); Throwable combined = CuVS2510GPUVectorsWriter.combineOperationAndCleanupFailures(bodyFailure, cleanupFailure); assertSame(bodyFailure, combined); - assertEquals(List.of("index", "dataset", "padded-dataset", "device-vectors"), closeOrder); + assertEquals(List.of("index", "dataset", "padded-dataset"), closeOrder); assertArrayEquals(new Throwable[] {indexFailure}, bodyFailure.getSuppressed()); assertArrayEquals( - new Throwable[] {datasetFailure, paddedFailure, deviceFailure}, - indexFailure.getSuppressed()); + new Throwable[] {datasetFailure, paddedFailure}, indexFailure.getSuppressed()); } private static AutoCloseable failingCloseable( @@ -147,4 +280,58 @@ private static AutoCloseable failingCloseable( throw failure; }; } + + private static CuVSDeviceMatrix fakeDeviceMatrix( + AtomicInteger closeCalls, RuntimeException closeFailure) { + return (CuVSDeviceMatrix) + Proxy.newProxyInstance( + CuVSDeviceMatrix.class.getClassLoader(), + new Class[] {CuVSDeviceMatrix.class}, + (proxy, method, args) -> { + if (method.getName().equals("close") && method.getParameterCount() == 0) { + closeCalls.incrementAndGet(); + if (closeFailure != null) { + throw closeFailure; + } + return null; + } + throw new AssertionError("Unexpected device-matrix call: " + method); + }); + } + + private static final class FakeDeviceBuilder implements CuVSMatrix.Builder { + private final CuVSDeviceMatrix dataset; + private final RuntimeException closeFailure; + private final AtomicInteger closeCalls = new AtomicInteger(); + + private FakeDeviceBuilder(CuVSDeviceMatrix dataset, RuntimeException closeFailure) { + this.dataset = dataset; + this.closeFailure = closeFailure; + } + + @Override + public void addVector(float[] vector) {} + + @Override + public void addVector(byte[] vector) {} + + @Override + public void addVector(int[] vector) {} + + @Override + public void addVector(short[] vector) {} + + @Override + public CuVSDeviceMatrix build() { + return dataset; + } + + @Override + public void close() { + closeCalls.incrementAndGet(); + if (closeFailure != null) { + throw closeFailure; + } + } + } } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java index df493c3e22..4ac27b988b 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java @@ -30,7 +30,9 @@ import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.NoMergePolicy; import org.apache.lucene.index.Term; +import org.apache.lucene.index.TieredMergePolicy; import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.store.ByteBuffersDirectory; import org.apache.lucene.store.Directory; @@ -80,7 +82,8 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { } try (Directory dir = newDirectory(new ByteBuffersDirectory()); - IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + IndexWriter w = + new IndexWriter(dir, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))) { for (int i = 0; i < R; i++) { Document doc = new Document(); doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); @@ -97,6 +100,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { assertEquals(1, subReaders.get(i).reader().getFloatVectorValues(F).size()); } } + w.getConfig().setMergePolicy(new TieredMergePolicy()); w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { @@ -177,7 +181,9 @@ public void testForceMergeUsesOnlyLiveSparseVectors() throws Exception { Map expected = new LinkedHashMap<>(); try (Directory directory = newDirectory(new ByteBuffersDirectory())) { - try (IndexWriter writer = new IndexWriter(directory, newIndexWriterConfig())) { + try (IndexWriter writer = + new IndexWriter( + directory, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))) { for (int segment = 0; segment < 3; segment++) { for (int row = 0; row < 5; row++) { String id = segment + "-" + row; @@ -198,6 +204,19 @@ public void testForceMergeUsesOnlyLiveSparseVectors() throws Exception { writer.deleteDocuments(new Term("id", segment + "-1")); } writer.commit(); + + try (DirectoryReader sourceReader = DirectoryReader.open(writer)) { + assertEquals("the test requires three source segments", 3, sourceReader.leaves().size()); + for (LeafReaderContext context : sourceReader.leaves()) { + LeafReader sourceLeaf = context.reader(); + assertTrue("each source segment must carry a deletion", sourceLeaf.hasDeletions()); + assertEquals(5, sourceLeaf.maxDoc()); + assertEquals(4, sourceLeaf.numDocs()); + assertEquals(4, sourceLeaf.getFloatVectorValues(vectorField).size()); + } + } + + writer.getConfig().setMergePolicy(new TieredMergePolicy()); writer.forceMerge(1); } @@ -243,7 +262,9 @@ private void assertTrivialLiveVectorMerge(int liveVectors) throws Exception { final String vectorField = "vector"; final int dimensions = 129; try (Directory directory = newDirectory(new ByteBuffersDirectory())) { - try (IndexWriter writer = new IndexWriter(directory, newIndexWriterConfig())) { + try (IndexWriter writer = + new IndexWriter( + directory, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))) { for (int id = 0; id < 3; id++) { Document vectorDocument = new Document(); vectorDocument.add(new StringField("id", "vector-" + id, Field.Store.YES)); @@ -260,6 +281,15 @@ private void assertTrivialLiveVectorMerge(int liveVectors) throws Exception { writer.deleteDocuments(new Term("id", "vector-" + id)); } writer.commit(); + + try (DirectoryReader sourceReader = DirectoryReader.open(writer)) { + assertEquals("the test requires three source segments", 3, sourceReader.leaves().size()); + for (LeafReaderContext context : sourceReader.leaves()) { + assertEquals(1, context.reader().getFloatVectorValues(vectorField).size()); + } + } + + writer.getConfig().setMergePolicy(new TieredMergePolicy()); writer.forceMerge(1); } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java index 6fdfa5a2a0..89011fc2ef 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java @@ -4,7 +4,10 @@ */ package com.nvidia.cuvs.lucene; +import com.nvidia.cuvs.CuVSMatrix; import java.io.IOException; +import java.lang.reflect.Proxy; +import java.util.concurrent.atomic.AtomicInteger; import org.apache.lucene.tests.util.LuceneTestCase; import org.junit.Test; @@ -48,6 +51,62 @@ public void testHandleThrowableWrapsCheckedExceptionWithCause() { assertSame(exception, thrown.getCause()); } + @Test + public void testMatrixBuilderCloseFailureClosesTransferredMatrix() { + IllegalStateException builderFailure = new IllegalStateException("builder close"); + AtomicInteger matrixCloses = new AtomicInteger(); + CuVSMatrix matrix = fakeMatrix(matrixCloses, null); + FakeMatrixBuilder builder = new FakeMatrixBuilder(matrix, builderFailure); + + IllegalStateException thrown = + assertThrows( + IllegalStateException.class, + () -> MatrixBuilderLifecycle.build(builder, CuVSMatrix.Builder::build)); + + assertSame(builderFailure, thrown); + assertEquals(1, builder.closeCalls.get()); + assertEquals(1, matrixCloses.get()); + } + + @Test + public void testMatrixCloseFailureIsSuppressedOnBuilderCloseFailure() { + IllegalStateException builderFailure = new IllegalStateException("builder close"); + IllegalArgumentException matrixFailure = new IllegalArgumentException("matrix close"); + AtomicInteger matrixCloses = new AtomicInteger(); + CuVSMatrix matrix = fakeMatrix(matrixCloses, matrixFailure); + FakeMatrixBuilder builder = new FakeMatrixBuilder(matrix, builderFailure); + + IllegalStateException thrown = + assertThrows( + IllegalStateException.class, + () -> MatrixBuilderLifecycle.build(builder, CuVSMatrix.Builder::build)); + + assertSame(builderFailure, thrown); + assertArrayEquals(new Throwable[] {matrixFailure}, thrown.getSuppressed()); + assertEquals(1, matrixCloses.get()); + } + + @Test + public void testMatrixOperationFailureRemainsPrimaryWhenBuilderCloseFails() { + IOException operationFailure = new IOException("populate"); + IllegalStateException builderFailure = new IllegalStateException("builder close"); + FakeMatrixBuilder builder = new FakeMatrixBuilder(null, builderFailure); + + IOException thrown = + assertThrows( + IOException.class, + () -> + MatrixBuilderLifecycle.build( + builder, + ignored -> { + throw operationFailure; + })); + + assertSame(operationFailure, thrown); + assertArrayEquals(new Throwable[] {builderFailure}, thrown.getSuppressed()); + assertEquals(1, builder.closeCalls.get()); + } + @Test public void testOwnedIndexClosesUntransferredDatasetOnce() throws Exception { TrackingCloseable dataset = new TrackingCloseable(null); @@ -116,4 +175,57 @@ public void close() throws Exception { } } } + + private static CuVSMatrix fakeMatrix(AtomicInteger closeCalls, RuntimeException closeFailure) { + return (CuVSMatrix) + Proxy.newProxyInstance( + CuVSMatrix.class.getClassLoader(), + new Class[] {CuVSMatrix.class}, + (proxy, method, args) -> { + if (method.getName().equals("close") && method.getParameterCount() == 0) { + closeCalls.incrementAndGet(); + if (closeFailure != null) { + throw closeFailure; + } + return null; + } + throw new AssertionError("Unexpected matrix call: " + method); + }); + } + + private static final class FakeMatrixBuilder implements CuVSMatrix.Builder { + private final CuVSMatrix matrix; + private final RuntimeException closeFailure; + private final AtomicInteger closeCalls = new AtomicInteger(); + + private FakeMatrixBuilder(CuVSMatrix matrix, RuntimeException closeFailure) { + this.matrix = matrix; + this.closeFailure = closeFailure; + } + + @Override + public void addVector(float[] vector) {} + + @Override + public void addVector(byte[] vector) {} + + @Override + public void addVector(int[] vector) {} + + @Override + public void addVector(short[] vector) {} + + @Override + public CuVSMatrix build() { + return matrix; + } + + @Override + public void close() { + closeCalls.incrementAndGet(); + if (closeFailure != null) { + throw closeFailure; + } + } + } } From b7d071706491151b1e9e829b794eac5f93f796ff Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 18 Sep 2026 18:18:00 +0000 Subject: [PATCH 06/21] Narrow accelerated HNSW memory fix scope Keep host-backed CAGRA-to-HNSW inputs, exact live-vector merge sizing, trivial merge handling, and the compatible upper-layer bridge. Restore the GPU-search codec to the target-branch device-input behavior and defer broader lifecycle, graph-integrity, and quantized-merge hardening to a follow-up. --- .../java-api-com-nvidia-cuvs-cuvsmatrix.md | 58 +- ...nvidia-cuvs-lucene-acceleratedhnswutils.md | 36 +- ...ia-cuvs-lucene-cuvs2510gpuvectorswriter.md | 45 +- ...om-nvidia-cuvs-lucene-gpubuilthnswgraph.md | 24 +- ...ne-lucene99acceleratedhnswvectorswriter.md | 16 +- ...leratedhnswbinaryquantizedvectorswriter.md | 10 +- ...leratedhnswscalarquantizedvectorswriter.md | 10 +- ...lucene-api-com-nvidia-cuvs-lucene-utils.md | 67 +- fern/pages/other/multidimensional_arrays.md | 12 +- .../com/nvidia/cuvs/CuVSDeviceMatrix.java | 18 +- .../java/com/nvidia/cuvs/CuVSHostMatrix.java | 19 +- .../main/java/com/nvidia/cuvs/CuVSMatrix.java | 19 +- .../nvidia/cuvs/MatrixBuilderLifecycle.java | 47 -- .../com/nvidia/cuvs/internal/common/Util.java | 5 +- .../com/nvidia/cuvs/spi/JDKProvider.java | 124 +--- .../java/com/nvidia/cuvs/CuVSMatrixIT.java | 34 - .../nvidia/cuvs/MatrixBuilderLifecycleIT.java | 137 ---- .../com/nvidia/cuvs/spi/CuVSProviderIT.java | 71 -- .../cuvs/lucene/AcceleratedHNSWUtils.java | 674 ++++++------------ .../cuvs/lucene/CuVS2510GPUVectorsWriter.java | 244 +------ .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 100 +-- .../Lucene99AcceleratedHNSWVectorsWriter.java | 57 +- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 120 +--- ...ratedHNSWScalarQuantizedVectorsWriter.java | 143 ++-- .../cuvs/lucene/MatrixBuilderLifecycle.java | 64 -- .../java/com/nvidia/cuvs/lucene/Utils.java | 202 +++--- .../TestAcceleratedHNSWDeletedDocuments.java | 6 - .../lucene/TestAcceleratedHNSWQuantizers.java | 55 -- .../TestAcceleratedHNSWUpperLayers.java | 228 +----- .../TestCagraIndexAtAlignedDimensions.java | 52 +- ...VS2510GPUVectorsWriterFailureHandling.java | 337 --------- .../lucene/TestQuantizedVectorsFormats.java | 153 +--- .../lucene/TestUtilsThrowableHandling.java | 181 ----- 33 files changed, 597 insertions(+), 2771 deletions(-) delete mode 100644 java/cuvs-java/src/main/java/com/nvidia/cuvs/MatrixBuilderLifecycle.java delete mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/MatrixBuilderLifecycleIT.java delete mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MatrixBuilderLifecycle.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWQuantizers.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsWriterFailureHandling.java diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsmatrix.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsmatrix.md index 0d62e838f4..29e3e618fd 100644 --- a/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsmatrix.md +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsmatrix.md @@ -114,34 +114,6 @@ Adds a single vector to the matrix. Each element is a raw float16 bit pattern st _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:102`_ -### build - -```java -T build() -``` - -Completes the matrix and transfers ownership to the caller. - -If this method fails, closing the builder releases any matrix storage allocated while the -builder was created. - -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:110`_ - -### close - -```java -@Override default void close() -``` - -Closes this builder. Built-in builders release matrix storage unless ownership was -transferred by a successful `#build()`. - -The default implementation preserves compatibility with providers compiled before -builders became closeable. Builders that allocate storage before `#build()` should -override this method. - -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:120`_ - ### hostBuilder ```java @@ -162,7 +134,7 @@ Returns a builder to create a new instance of a host-memory matrix a builder for creating a `CuVSHostMatrix` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:132`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:115`_ ### hostBuilder @@ -186,7 +158,7 @@ Returns a builder to create a new instance of a host-memory matrix a builder for creating a `CuVSDeviceMatrix` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:146`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:129`_ ### deviceBuilder @@ -209,7 +181,7 @@ Returns a builder to create a new instance of a dataset a builder for creating a `CuVSDeviceMatrix` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:161`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:144`_ ### deviceBuilder @@ -234,7 +206,7 @@ Returns a builder to create a new instance of a dataset a builder for creating a `CuVSDeviceMatrix` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:177`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:160`_ ### size @@ -248,7 +220,7 @@ Gets the size of the dataset Size of the dataset -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:193`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:176`_ ### columns @@ -263,7 +235,7 @@ or the graph degree for the graph represented as a list of neighbours Dimensions of the vectors in the dataset -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:201`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:184`_ ### dataType @@ -277,7 +249,7 @@ Gets the element type a `DataType` describing the matrix element type -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:208`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:191`_ ### getRow @@ -293,7 +265,7 @@ Get a view (0-copy) of the row data, as a list of integers (32 bit) | --- | --- | | `row` | the row for which to return the data | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:215`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:198`_ ### toArray @@ -309,7 +281,7 @@ Copies the content of this dataset to an on-heap Java matrix (array of arrays). | --- | --- | | `array` | the destination array. Must be of length `CuVSMatrix#size()` or bigger, and each element must be of length `CuVSMatrix#columns()` or bigger. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:223`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:206`_ ### toArray @@ -325,7 +297,7 @@ Copies the content of this dataset to an on-heap Java matrix (array of arrays). | --- | --- | | `array` | the destination array. Must be of length `CuVSMatrix#size()` or bigger, and each element must be of length `CuVSMatrix#columns()` or bigger. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:231`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:214`_ ### toArray @@ -341,7 +313,7 @@ Copies the content of this dataset to an on-heap Java matrix (array of arrays). | --- | --- | | `array` | the destination array. Must be of length `CuVSMatrix#size()` or bigger, and each element must be of length `CuVSMatrix#columns()` or bigger. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:239`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:222`_ ### toHost @@ -359,7 +331,7 @@ same element type and dimension. | --- | --- | | `hostMatrix` | the host-memory-backed matrix to fill. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:248`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:231`_ ### toHost @@ -373,7 +345,7 @@ the device matrix. The returned host matrix will need to be managed by the caller, which will be responsible to call `CuVSMatrix#close()` to free its resources when done. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:257`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:240`_ ### toDevice @@ -391,7 +363,7 @@ same element type and dimension. | --- | --- | | `deviceMatrix` | the device-memory-backed matrix to fill. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:266`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:249`_ ### toDevice @@ -405,6 +377,6 @@ the host matrix. The returned device matrix will need to be managed by the caller, which will be responsible to call `CuVSMatrix#close()` to free its resources when done. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:275`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:258`_ _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:17`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md index 0aa315d910..219149bc6d 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md @@ -29,25 +29,13 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWU public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, int size, int dimensions, CuVSMatrix adjacencyListMatrix, List vectors, int hnswLayers, CagraIndexParams params, QuantizationType quantization) throws Throwable ``` -Creates a multi-layer HNSW graph from heap vectors. +Creates a multi-layer HNSW graph with dynamic number of layers. +M = ceil(cagraGraphDegree / 2), where cagraGraphDegree is the CAGRA adjacency list's degree +(its column count). Ceil is used to accommodate odd graph degrees. +Each layer contains 1/M nodes from the previous layer +Creates layers until the highest layer has ≤ M nodes -This overload preserves the original public API. Only rows selected for an upper layer are -copied into native memory. - -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:74`_ - -### createMultiLayerHnswGraph - -```java -static GPUBuiltHnswGraph createMultiLayerHnswGraph( int dimensions, CuVSMatrix adjacencyListMatrix, CuVSMatrix vectorDataset, int hnswLayers, CagraIndexParams params, QuantizationType quantization) throws Throwable -``` - -Creates a multi-layer HNSW graph from a native matrix. - -Only sampled rows are copied into each upper-layer matrix; the complete dataset is never -materialized on the Java heap. - -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:106`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:81`_ ### writeGraph @@ -74,7 +62,7 @@ a 2D array of offsets | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:443`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:237`_ ### writeMeta @@ -103,7 +91,7 @@ Writes the meta information for the index. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:508`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:302`_ ### printInfoStream @@ -119,7 +107,7 @@ A utility method to print info/debugging messages using InfoStream. | --- | --- | | `msg` | the debugging message to print | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:590`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:384`_ ### writeEmpty @@ -141,7 +129,7 @@ Writes an empty meta information for the field. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:602`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:396`_ ### quantizeFloatVectorsToBinary @@ -164,7 +152,7 @@ Bits are packed: 8 dimensions per byte. A list of byte binary representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:615`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:409`_ ### quantizeFloatVectorsToScalar @@ -184,6 +172,6 @@ Scalar quantization. A list of byte scalar representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:643`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:451`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:31`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorswriter.md index 956074833f..df1abbbdc7 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorswriter.md @@ -22,7 +22,7 @@ CAGRA(true, false), /** Builds a Brute Force index. */ BRUTE_FORCE(false, true), Builds a CAGRA index. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:91`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:89`_ ### BRUTE_FORCE @@ -32,7 +32,7 @@ BRUTE_FORCE(false, true), /** Builds both - CAGRA and Brute Force indexes. */ CA Builds a Brute Force index. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:94`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:92`_ ### CAGRA_AND_BRUTE_FORCE @@ -42,7 +42,7 @@ CAGRA_AND_BRUTE_FORCE(true, true) Builds both - CAGRA and Brute Force indexes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:97`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:95`_ ### CuVS2510GPUVectorsWriter @@ -66,7 +66,7 @@ Initializes `CuVS2510GPUVectorsWriter`. | --- | --- | | `IOException` | I/O exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:123`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:121`_ ### addField @@ -76,30 +76,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:162`_ - -### prepareCagraDataset - -```java -static CuVSDeviceMatrix prepareCagraDataset( CagraDatasetBuilderFactory builderFactory, CagraDatasetPopulation population) throws Throwable -``` - -Builds the device input and applies the pre-persistence fallback policy only after builder -cleanup succeeds. A builder cleanup failure is fatal; after ownership has transferred, this -method also closes the dataset before propagating that failure. A builder-factory failure is -also fatal because no cleanup handle was returned. - -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:276`_ - -### closeCagraResources - -```java -static Throwable closeCagraResources( AutoCloseable index, AutoCloseable originalDataset, AutoCloseable indexDataset) -``` - -Closes CAGRA resources in dependency order and returns the first cleanup failure. - -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:391`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:160`_ ### flush @@ -109,7 +86,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Creates the CAGRA and/or brute force indexes and writes them to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:471`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:303`_ ### mergeOneField @@ -119,7 +96,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:853`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:685`_ ### ramBytesUsed @@ -129,7 +106,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:864`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:696`_ ### finish @@ -139,7 +116,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:876`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:708`_ ### close @@ -149,6 +126,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Close the applicable resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:896`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:728`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:61`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java:59`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md index 91ed1f3a3d..f8a3aee9dc 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md @@ -31,7 +31,7 @@ Multi-layer constructor that supports arbitrary number of layers. | `layerNodes` | the nodes on the layer | | `layerAdjacencies` | adjacency list | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:44`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:41`_ ### getNodesOnLevel @@ -41,7 +41,7 @@ public NodesIterator getNodesOnLevel(int level) Get all nodes on a given level as node 0th ordinals. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:161`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:89`_ ### getNeighbors @@ -62,7 +62,7 @@ Get the neighbors for the node and the level it resides. an instance of NeighborArray -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:179`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:107`_ ### seek @@ -72,7 +72,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Move the pointer to exactly the given level's target. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:206`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:132`_ ### nextNeighbor @@ -82,7 +82,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Iterates over the neighbor list. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:216`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:142`_ ### entryNode @@ -92,7 +92,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns graph's entry point on the top level. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:247`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:173`_ ### maxConn @@ -102,7 +102,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap returns M, the maximum number of connections for a node. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:266`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:192`_ ### neighborCount @@ -112,7 +112,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns the neighbor count. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:281`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:207`_ ### size @@ -122,7 +122,7 @@ public int size() Returns the number of nodes in the graph. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:356`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:282`_ ### numLevels @@ -136,7 +136,7 @@ Returns the number of levels in the HNSW graph. the number of levels -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:365`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:291`_ ### dimensions @@ -150,6 +150,6 @@ Gets the vector dimension. the vector dimension -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:374`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:300`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:23`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:21`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md index 0bd4d8b8b3..2f30caeb07 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md @@ -37,7 +37,7 @@ Initializes `Lucene99AcceleratedHNSWVectorsWriter` | --- | --- | | `IOException` | IOException | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:89`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:86`_ ### addField @@ -47,7 +47,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:131`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:128`_ ### flush @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:226`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:203`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:339`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:291`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:348`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:300`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:368`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:320`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:378`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:330`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:55`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:52`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md index 98b9ee0aa4..91a4e67303 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:227`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:215`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:312`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:302`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:383`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:333`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:403`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:353`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:412`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:362`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:57`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md index 04f5957faa..9c51cb26e9 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:244`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:241`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:327`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:326`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:404`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:357`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:424`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:377`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:433`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:386`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:56`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md index be145bc37a..0d2a9cb1a6 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md @@ -50,34 +50,38 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:43`_ ### createFloatMatrix ```java -static CuVSMatrix createFloatMatrix(List data, int dimensions) throws IOException +static CuVSMatrix createFloatMatrix(List data, int dimensions, CuVSResources resources) ``` -Builds a host-memory CuVSMatrix from a list of float vectors. +A method to build a CuVSMatrix from a list of float vectors. -Copies vectors directly into a native host matrix via `CuVSMatrix#hostBuilder`, -without creating an intermediate `float[][]` on the heap. +Uses CuVSMatrix.Builder to copy vectors directly to device memory +without creating intermediate heap arrays. **Parameters** | Name | Description | | --- | --- | | `data` | The float vectors | -| `dimensions` | The number of float elements in each vector | +| `dimensions` | The number float elements in each vector | +| `resources` | The CuVS resources for device matrix creation | **Returns** -a host-memory CuVSMatrix +an instance of CuVSMatrix -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:62`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:63`_ ### createByteMatrix ```java -static CuVSMatrix createByteMatrix(List data, int bytesPerVector) throws IOException +static CuVSMatrix createByteMatrix( List data, int bytesPerVector, CuVSResources resources) ``` -Builds a host-memory CuVSMatrix from a list of byte vectors (e.g. quantized vectors). +A method to build a CuVSMatrix from a list of byte vectors (for binary quantized vectors). + +Uses CuVSMatrix.Builder to copy vectors directly to device memory +without creating intermediate heap arrays. **Parameters** @@ -85,20 +89,21 @@ Builds a host-memory CuVSMatrix from a list of byte vectors (e.g. quantized vect | --- | --- | | `data` | The byte vectors (packed bits for binary quantization) | | `bytesPerVector` | The number of bytes in each vector | +| `resources` | The CuVS resources for device matrix creation | **Returns** -a host-memory CuVSMatrix with BYTE data type +an instance of CuVSMatrix with BYTE data type -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:80`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:92`_ ### createByteMatrixFromArray ```java -static CuVSMatrix createByteMatrixFromArray(byte[][] data, int bytesPerVector) throws IOException +static CuVSMatrix createByteMatrixFromArray( byte[][] data, int bytesPerVector, CuVSResources resources) ``` -Builds a host-memory CuVSMatrix from a 2D byte array (e.g. quantized vectors). +A method to build a CuVSMatrix from a 2D byte array (for binary quantized vectors). **Parameters** @@ -106,33 +111,13 @@ Builds a host-memory CuVSMatrix from a 2D byte array (e.g. quantized vectors). | --- | --- | | `data` | The 2D byte array (packed bits for binary quantization) | | `bytesPerVector` | The number of bytes in each vector | +| `resources` | The CuVS resources for device matrix creation | **Returns** -a host-memory CuVSMatrix with BYTE data type - -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:98`_ - -### closeIndexWithDatasetFallback - -```java -static void closeIndexWithDatasetFallback(AutoCloseable index, AutoCloseable dataset) throws Exception -``` - -Closes an index that owns `dataset`. If index cleanup fails before releasing the -dataset, a direct dataset close is attempted and attached to the index failure when needed. - -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:114`_ - -### ownDataset - -```java -static OwnedIndex ownDataset(AutoCloseable dataset) -``` - -Starts an ownership scope for a dataset that may later be transferred to an index. +an instance of CuVSMatrix with BYTE data type -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:131`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:119`_ ### nanosToMillis @@ -152,7 +137,7 @@ A utility method to convert nanoseconds to milliseconds. milliseconds -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:196`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:141`_ ### cuVSResourcesOrNull @@ -166,7 +151,7 @@ Creates an instance of CuVSResources. an instance of CuVSResources -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:205`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:150`_ ### handleThrowableWithIgnore @@ -189,7 +174,7 @@ A utility method that conditionally ignores certain throwable objects | --- | --- | | `IOException` | | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:233`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:178`_ ### createListFromMergedVectors @@ -215,7 +200,7 @@ a list of float arrays | --- | --- | | `IOException` | I/O Exception | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:247`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:192`_ ### info @@ -233,6 +218,6 @@ Utility to print info/debug messages via InfoStream. | `component` | the name of the index writer | | `msg` | the log message to push via the InfoStream | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:265`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:210`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:26`_ diff --git a/fern/pages/other/multidimensional_arrays.md b/fern/pages/other/multidimensional_arrays.md index 072afa1087..c5f7ec7d5e 100644 --- a/fern/pages/other/multidimensional_arrays.md +++ b/fern/pages/other/multidimensional_arrays.md @@ -165,16 +165,16 @@ import com.nvidia.cuvs.CuVSDeviceMatrix; long nRows = 100_000; long nFeatures = 128; -try (CuVSResources resources = CuVSResources.create(); - CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, nRows, nFeatures, CuVSMatrix.DataType.FLOAT)) { +try (CuVSResources resources = CuVSResources.create()) { + CuVSMatrix.Builder builder = + CuVSMatrix.deviceBuilder( + resources, nRows, nFeatures, CuVSMatrix.DataType.FLOAT); + for (long row = 0; row < nRows; row++) { builder.addVector(loadVector(row)); } - // A successful build transfers matrix ownership from the builder. - try (CuVSDeviceMatrix dataset = builder.build()) { + try (CuVSMatrix dataset = builder.build()) { // Pass dataset to NVIDIA cuVS Java APIs. } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java index e0065cd91b..16b1881ad5 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSDeviceMatrix.java @@ -15,20 +15,8 @@ public interface CuVSDeviceMatrix extends CuVSMatrix { * responsible to call {@link CuVSMatrix#close()} to free its resources when done. */ default CuVSHostMatrix toHost() { - CuVSHostMatrix hostMatrix = - MatrixBuilderLifecycle.buildAndClose(CuVSMatrix.hostBuilder(size(), columns(), dataType())); - try { - toHost(hostMatrix); - return hostMatrix; - } catch (RuntimeException | Error failure) { - try { - hostMatrix.close(); - } catch (RuntimeException | Error closeFailure) { - if (failure != closeFailure) { - failure.addSuppressed(closeFailure); - } - } - throw failure; - } + var hostMatrix = CuVSMatrix.hostBuilder(size(), columns(), dataType()).build(); + toHost(hostMatrix); + return hostMatrix; } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java index e8014d73e4..62b708f71a 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSHostMatrix.java @@ -11,21 +11,8 @@ public interface CuVSHostMatrix extends CuVSMatrix { int get(int row, int col); default CuVSDeviceMatrix toDevice(CuVSResources resources) { - CuVSDeviceMatrix deviceMatrix = - MatrixBuilderLifecycle.buildAndClose( - CuVSMatrix.deviceBuilder(resources, size(), columns(), dataType())); - try { - toDevice(deviceMatrix, resources); - return deviceMatrix; - } catch (RuntimeException | Error failure) { - try { - deviceMatrix.close(); - } catch (RuntimeException | Error closeFailure) { - if (failure != closeFailure) { - failure.addSuppressed(closeFailure); - } - } - throw failure; - } + var deviceMatrix = CuVSMatrix.deviceBuilder(resources, size(), columns(), dataType()).build(); + toDevice(deviceMatrix, resources); + return deviceMatrix; } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java index 4447f88db6..097efb1003 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java @@ -73,7 +73,7 @@ static CuVSMatrix ofArray(byte[][] vectors) { * A builder to construct a new matrix one row at a time * @param the CuVSMatrix type to build */ - interface Builder extends AutoCloseable { + interface Builder { /** * Adds a single vector to the matrix. * @@ -102,24 +102,7 @@ interface Builder extends AutoCloseable { */ void addVector(short[] vector); - /** - * Completes the matrix and transfers ownership to the caller. - * - *

If this method fails, closing the builder releases any matrix storage allocated while the - * builder was created. - */ T build(); - - /** - * Closes this builder. Built-in builders release matrix storage unless ownership was - * transferred by a successful {@link #build()}. - * - *

The default implementation preserves compatibility with providers compiled before - * builders became closeable. Builders that allocate storage before {@link #build()} should - * override this method. - */ - @Override - default void close() {} } /** diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/MatrixBuilderLifecycle.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/MatrixBuilderLifecycle.java deleted file mode 100644 index c858535039..0000000000 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/MatrixBuilderLifecycle.java +++ /dev/null @@ -1,47 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs; - -import java.util.Objects; - -/** Ownership-aware finalization shared by the public matrix conversion defaults. */ -final class MatrixBuilderLifecycle { - - private MatrixBuilderLifecycle() {} - - static T buildAndClose(CuVSMatrix.Builder builder) { - T matrix; - try { - matrix = Objects.requireNonNull(builder.build(), "Matrix builder must not return null"); - } catch (RuntimeException | Error operationFailure) { - closeAndSuppress(builder, operationFailure); - throw operationFailure; - } - - try { - builder.close(); - return matrix; - } catch (RuntimeException | Error builderCloseFailure) { - closeAndSuppress(matrix, builderCloseFailure); - throw builderCloseFailure; - } - } - - private static void closeAndSuppress(AutoCloseable resource, Throwable failure) { - try { - resource.close(); - } catch (RuntimeException | Error closeFailure) { - if (failure != closeFailure) { - failure.addSuppressed(closeFailure); - } - } catch (Exception closeFailure) { - // CuVSMatrix and its Builder narrow close() to unchecked failures. Keep this guard so an - // unusual AutoCloseable implementation cannot replace the original failure. - if (failure != closeFailure) { - failure.addSuppressed(closeFailure); - } - } - } -} diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java index 8025e9c640..986ca5bd88 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/common/Util.java @@ -11,7 +11,6 @@ import static com.nvidia.cuvs.internal.panama.headers_h.*; import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.LibraryException; import com.nvidia.cuvs.internal.panama.DLDataType; import com.nvidia.cuvs.internal.panama.DLDevice; import com.nvidia.cuvs.internal.panama.DLManagedTensor; @@ -127,7 +126,7 @@ public static int cudaGetDeviceProperties(MemorySegment prop, int device) { public static void checkCuVSError(int value, String caller) { if (value != CUVS_SUCCESS) { String errorMsg = getLastErrorText(); - throw new LibraryException(caller + " returned " + value + "[" + errorMsg + "]"); + throw new RuntimeException(caller + " returned " + value + "[" + errorMsg + "]"); } } @@ -139,7 +138,7 @@ public static void checkCuVSError(int value, String caller) { */ public static void checkCudaError(int value, String caller) { if (value != CUDA_SUCCESS) { - throw new LibraryException(caller + " returned " + value); + throw new RuntimeException(caller + " returned " + value); } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index eca0c96336..1f16a4e904 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -30,8 +30,6 @@ import java.util.List; import java.util.Locale; import java.util.Objects; -import java.util.function.Function; -import java.util.function.Supplier; import java.util.jar.JarFile; import java.util.jar.Manifest; import java.util.logging.Level; @@ -546,15 +544,10 @@ public CuVSMatrix.Builder newHostMatrixBuilder( @Override public CuVSMatrix.Builder newDeviceMatrixBuilder( CuVSResources resources, long size, long columns, CuVSMatrix.DataType dataType) { - var rowBytes = deviceMatrixRowBytes(size, columns, dataType); - var bufferRowCount = Math.min(PinnedMemoryBuffer.CHUNK_BYTES / rowBytes, size); - return createDeviceMatrixBuilder( - () -> Util.getStream(resources), - stream -> - rowBytes > PinnedMemoryBuffer.CHUNK_BYTES - ? new DirectDeviceMatrixBuilder(resources, size, columns, dataType, stream) - : new BufferedDeviceMatrixBuilder( - resources, size, columns, dataType, stream, bufferRowCount)); + var rowBytes = columns * dataType.bytes(); + return rowBytes > PinnedMemoryBuffer.CHUNK_BYTES + ? new DirectDeviceMatrixBuilder(resources, size, columns, dataType) + : new BufferedDeviceMatrixBuilder(resources, size, columns, dataType); } @Override @@ -565,40 +558,11 @@ public CuVSMatrix.Builder newDeviceMatrixBuilder( int rowStride, int columnStride, CuVSMatrix.DataType dataType) { - var rowBytes = deviceMatrixRowBytes(size, columns, dataType); - var bufferRowCount = Math.min(PinnedMemoryBuffer.CHUNK_BYTES / rowBytes, size); - return createDeviceMatrixBuilder( - () -> Util.getStream(resources), - stream -> - rowBytes > PinnedMemoryBuffer.CHUNK_BYTES - ? new DirectDeviceMatrixBuilder( - resources, size, columns, rowStride, columnStride, dataType, stream) - : new BufferedDeviceMatrixBuilder( - resources, - size, - columns, - rowStride, - columnStride, - dataType, - stream, - bufferRowCount)); - } - - private static long deviceMatrixRowBytes(long size, long columns, CuVSMatrix.DataType dataType) { - if (size < 0) { - throw new IllegalArgumentException("size must be non-negative: " + size); - } - if (columns <= 0) { - throw new IllegalArgumentException("columns must be positive: " + columns); - } - return Math.multiplyExact(columns, dataType.bytes()); - } - - /** Acquires the stream before invoking a constructor that allocates device memory. */ - static T createDeviceMatrixBuilder( - Supplier streamSupplier, Function builderFactory) { - MemorySegment stream = streamSupplier.get(); - return builderFactory.apply(stream); + var rowBytes = columns * dataType.bytes(); + return rowBytes > PinnedMemoryBuffer.CHUNK_BYTES + ? new DirectDeviceMatrixBuilder(resources, size, columns, rowStride, columnStride, dataType) + : new BufferedDeviceMatrixBuilder( + resources, size, columns, rowStride, columnStride, dataType); } @Override @@ -653,8 +617,7 @@ public CuVSMatrix newMatrixFromArray(byte[][] vectors) { return dataset; } - private abstract static class MatrixBuilder - implements AutoCloseable { + private abstract static class MatrixBuilder { protected final long columns; protected final long size; @@ -663,7 +626,6 @@ private abstract static class MatrixBuilder protected final long rowSize; protected final long rowBytes; protected int currentRow; - private boolean closed; protected MatrixBuilder(T matrix, long size, long columns) { this.columns = columns; @@ -673,7 +635,6 @@ protected MatrixBuilder(T matrix, long size, long columns) { this.rowSize = columns * elementSize; this.rowBytes = rowSize; this.currentRow = 0; - this.closed = false; } protected MatrixBuilder(T matrix, long size, long columns, int rowStride) { @@ -685,11 +646,9 @@ protected MatrixBuilder(T matrix, long size, long columns, int rowStride) { this.rowBytes = columns * elementSize; this.currentRow = 0; - this.closed = false; } public void addVector(float[] vector) { - ensureOpen(); if (vector.length != columns) { throw new IllegalArgumentException( String.format( @@ -699,7 +658,6 @@ public void addVector(float[] vector) { } public void addVector(byte[] vector) { - ensureOpen(); if (vector.length != columns) { throw new IllegalArgumentException( String.format( @@ -709,7 +667,6 @@ public void addVector(byte[] vector) { } public void addVector(int[] vector) { - ensureOpen(); if (vector.length != columns) { throw new IllegalArgumentException( String.format( @@ -719,7 +676,6 @@ public void addVector(int[] vector) { } public void addVector(short[] vector) { - ensureOpen(); if (vector.length != columns) { throw new IllegalArgumentException( String.format( @@ -728,27 +684,6 @@ public void addVector(short[] vector) { internalAddVector(MemorySegment.ofArray(vector)); } - protected final T transferOwnership() { - ensureOpen(); - closed = true; - return matrix; - } - - protected final void ensureOpen() { - if (closed) { - throw new IllegalStateException("matrix builder is closed"); - } - } - - @Override - public final void close() { - if (closed) { - return; - } - closed = true; - matrix.close(); - } - protected abstract void internalAddVector(MemorySegment vector); } @@ -766,17 +701,12 @@ private static final class BufferedDeviceMatrixBuilder extends MatrixBuilder builder.addVector(new float[4])); - } - } - - @Test - public void testClosingDeviceBuilderAfterBuildDoesNotCloseMatrix() throws Throwable { - try (var resources = CheckedCuVSResources.create()) { - CuVSDeviceMatrix matrix; - try (var builder = CuVSMatrix.deviceBuilder(resources, 1, 4, CuVSMatrix.DataType.FLOAT)) { - builder.addVector(new float[] {1f, 2f, 3f, 4f}); - matrix = builder.build(); - - assertThrows(IllegalStateException.class, builder::build); - assertThrows(IllegalStateException.class, () -> builder.addVector(new float[4])); - } - - try (matrix) { - float[][] actual = new float[1][4]; - matrix.toArray(actual); - assertArrayEquals(new float[] {1f, 2f, 3f, 4f}, actual[0], DELTA); - } - } - } - private void testIntDatasetBuilder(int rows, int cols, CuVSMatrix.Builder builder) { var data = new int[rows][cols]; diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/MatrixBuilderLifecycleIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/MatrixBuilderLifecycleIT.java deleted file mode 100644 index 1fe35025f8..0000000000 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/MatrixBuilderLifecycleIT.java +++ /dev/null @@ -1,137 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs; - -import static org.junit.Assert.assertArrayEquals; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertSame; -import static org.junit.Assert.assertThrows; - -import java.lang.reflect.Proxy; -import java.util.concurrent.atomic.AtomicInteger; -import org.junit.Test; - -public class MatrixBuilderLifecycleIT { - - @Test - public void testSuccessfulBuildTransfersOwnershipAfterBuilderClose() { - AtomicInteger matrixCloses = new AtomicInteger(); - CuVSMatrix matrix = fakeMatrix(matrixCloses, null); - FakeBuilder builder = new FakeBuilder(matrix, null, null); - - CuVSMatrix built = MatrixBuilderLifecycle.buildAndClose(builder); - - assertSame(matrix, built); - assertEquals(1, builder.closeCalls.get()); - assertEquals(0, matrixCloses.get()); - } - - @Test - public void testBuilderCloseFailureClosesTransferredMatrix() { - IllegalStateException builderFailure = new IllegalStateException("builder close"); - AtomicInteger matrixCloses = new AtomicInteger(); - CuVSMatrix matrix = fakeMatrix(matrixCloses, null); - FakeBuilder builder = new FakeBuilder(matrix, null, builderFailure); - - IllegalStateException thrown = - assertThrows( - IllegalStateException.class, () -> MatrixBuilderLifecycle.buildAndClose(builder)); - - assertSame(builderFailure, thrown); - assertEquals(1, builder.closeCalls.get()); - assertEquals(1, matrixCloses.get()); - } - - @Test - public void testMatrixCloseFailureIsSuppressedOnBuilderCloseFailure() { - IllegalStateException builderFailure = new IllegalStateException("builder close"); - IllegalArgumentException matrixFailure = new IllegalArgumentException("matrix close"); - AtomicInteger matrixCloses = new AtomicInteger(); - CuVSMatrix matrix = fakeMatrix(matrixCloses, matrixFailure); - FakeBuilder builder = new FakeBuilder(matrix, null, builderFailure); - - IllegalStateException thrown = - assertThrows( - IllegalStateException.class, () -> MatrixBuilderLifecycle.buildAndClose(builder)); - - assertSame(builderFailure, thrown); - assertArrayEquals(new Throwable[] {matrixFailure}, thrown.getSuppressed()); - assertEquals(1, matrixCloses.get()); - } - - @Test - public void testBuildFailureRemainsPrimaryWhenBuilderCloseAlsoFails() { - IllegalArgumentException buildFailure = new IllegalArgumentException("build"); - IllegalStateException closeFailure = new IllegalStateException("builder close"); - FakeBuilder builder = new FakeBuilder(null, buildFailure, closeFailure); - - IllegalArgumentException thrown = - assertThrows( - IllegalArgumentException.class, () -> MatrixBuilderLifecycle.buildAndClose(builder)); - - assertSame(buildFailure, thrown); - assertArrayEquals(new Throwable[] {closeFailure}, thrown.getSuppressed()); - assertEquals(1, builder.closeCalls.get()); - } - - private static CuVSMatrix fakeMatrix(AtomicInteger closeCalls, RuntimeException closeFailure) { - return (CuVSMatrix) - Proxy.newProxyInstance( - CuVSMatrix.class.getClassLoader(), - new Class[] {CuVSMatrix.class}, - (proxy, method, args) -> { - if (method.getName().equals("close") && method.getParameterCount() == 0) { - closeCalls.incrementAndGet(); - if (closeFailure != null) { - throw closeFailure; - } - return null; - } - throw new AssertionError("Unexpected matrix call: " + method); - }); - } - - private static final class FakeBuilder implements CuVSMatrix.Builder { - private final CuVSMatrix matrix; - private final RuntimeException buildFailure; - private final RuntimeException closeFailure; - private final AtomicInteger closeCalls = new AtomicInteger(); - - private FakeBuilder( - CuVSMatrix matrix, RuntimeException buildFailure, RuntimeException closeFailure) { - this.matrix = matrix; - this.buildFailure = buildFailure; - this.closeFailure = closeFailure; - } - - @Override - public void addVector(float[] vector) {} - - @Override - public void addVector(byte[] vector) {} - - @Override - public void addVector(int[] vector) {} - - @Override - public void addVector(short[] vector) {} - - @Override - public CuVSMatrix build() { - if (buildFailure != null) { - throw buildFailure; - } - return matrix; - } - - @Override - public void close() { - closeCalls.incrementAndGet(); - if (closeFailure != null) { - throw closeFailure; - } - } - } -} diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java index b9ed704d70..076db40553 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/spi/CuVSProviderIT.java @@ -6,18 +6,11 @@ import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; import com.nvidia.cuvs.CuVSTestCase; -import com.nvidia.cuvs.LibraryException; -import java.lang.foreign.MemorySegment; import java.lang.invoke.MethodHandles; import java.lang.invoke.MethodType; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Function; -import java.util.function.Supplier; import org.junit.Before; import org.junit.Test; @@ -125,57 +118,6 @@ public void testMaxVersionOverride() { } } - @Test - public void testDeviceMatrixBuilderDoesNotAllocateWhenStreamAcquisitionFails() { - LibraryException streamFailure = new LibraryException("stream acquisition"); - AtomicInteger streamCalls = new AtomicInteger(); - AtomicInteger builderCalls = new AtomicInteger(); - - LibraryException thrown = - assertThrows( - LibraryException.class, - () -> - createDeviceMatrixBuilder( - () -> { - streamCalls.incrementAndGet(); - throw streamFailure; - }, - ignored -> { - builderCalls.incrementAndGet(); - return new Object(); - })); - - assertSame(streamFailure, thrown); - assertEquals(1, streamCalls.get()); - assertEquals(0, builderCalls.get()); - } - - @Test - public void testDeviceMatrixBuilderPassesAcquiredStreamToConstructor() throws Throwable { - MemorySegment stream = MemorySegment.ofArray(new byte[1]); - Object builder = new Object(); - AtomicInteger streamCalls = new AtomicInteger(); - AtomicInteger builderCalls = new AtomicInteger(); - AtomicReference observedStream = new AtomicReference<>(); - - Object created = - createDeviceMatrixBuilder( - () -> { - streamCalls.incrementAndGet(); - return stream; - }, - actualStream -> { - builderCalls.incrementAndGet(); - observedStream.set(actualStream); - return builder; - }); - - assertSame(builder, created); - assertSame(stream, observedStream.get()); - assertEquals(1, streamCalls.get()); - assertEquals(1, builderCalls.get()); - } - static void checkCuVSVersionMatching(String mavenVersionString, int major, int minor, int patch) throws ProviderInitializationException { try { @@ -194,17 +136,4 @@ static void checkCuVSVersionMatching(String mavenVersionString, int major, int m throw new AssertionError(e); } } - - static Object createDeviceMatrixBuilder( - Supplier streamSupplier, Function builderFactory) - throws Throwable { - var cls = Class.forName("com.nvidia.cuvs.spi.JDKProvider"); - var method = - MethodHandles.lookup() - .findStatic( - cls, - "createDeviceMatrixBuilder", - MethodType.methodType(Object.class, Supplier.class, Function.class)); - return method.invoke(streamSupplier, builderFactory); - } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index a4c96b9df6..09325bc794 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -6,16 +6,17 @@ package com.nvidia.cuvs.lucene; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.Utils.createHostByteMatrixFromArray; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.RowView; import java.io.IOException; +import java.util.AbstractList; import java.util.ArrayList; import java.util.Arrays; import java.util.List; -import java.util.Objects; import java.util.Random; import java.util.SortedSet; import java.util.TreeSet; @@ -58,19 +59,26 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens int[][] singleNodeAdjacency = new int[][] {{-1}}; // -1 indicates no neighbors // Create CuVSMatrix from the adjacency list - try (CuVSMatrix adjacencyMatrix = CuVSMatrix.ofArray(singleNodeAdjacency)) { - // Layer 0 contains every node, so its node list is implicit. - List layerNodes = new ArrayList<>(); - layerNodes.add(null); - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, List.of(adjacencyMatrix)); - } + CuVSMatrix adjacencyMatrix = CuVSMatrix.ofArray(singleNodeAdjacency); + + // Create layer data for single-level graph + List layerNodes = new ArrayList<>(); + List layerAdjacencies = new ArrayList<>(); + + // Layer 0: contains all nodes (just the single node) + layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list + layerAdjacencies.add(adjacencyMatrix); + + // Create the single-layer graph + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); } /** - * Creates a multi-layer HNSW graph from heap vectors. - * - *

This overload preserves the original public API. Only rows selected for an upper layer are - * copied into native memory. + * Creates a multi-layer HNSW graph with dynamic number of layers. + * M = ceil(cagraGraphDegree / 2), where cagraGraphDegree is the CAGRA adjacency list's degree + * (its column count). Ceil is used to accommodate odd graph degrees. + * Each layer contains 1/M nodes from the previous layer + * Creates layers until the highest layer has ≤ M nodes */ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, @@ -82,355 +90,192 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( CagraIndexParams params, QuantizationType quantization) throws Throwable { - Objects.requireNonNull(vectors, "vectors"); - if (vectors.size() != size) { - throw new IllegalArgumentException( - "Expected " + size + " vectors, but received " + vectors.size()); - } - CuVSMatrix.DataType dataType = expectedDataType(quantization); - int columns = expectedColumns(dimensions, quantization); - return createMultiLayerHnswGraph( - size, - dimensions, - adjacencyListMatrix, - hnswLayers, - params, - selectedNodes -> createSubsetDataset(vectors, selectedNodes, columns, dataType)); - } - /** - * Creates a multi-layer HNSW graph from a native matrix. - * - *

Only sampled rows are copied into each upper-layer matrix; the complete dataset is never - * materialized on the Java heap. - */ - static GPUBuiltHnswGraph createMultiLayerHnswGraph( - int dimensions, - CuVSMatrix adjacencyListMatrix, - CuVSMatrix vectorDataset, - int hnswLayers, - CagraIndexParams params, - QuantizationType quantization) - throws Throwable { - Objects.requireNonNull(vectorDataset, "vectorDataset"); - int size = Math.toIntExact(vectorDataset.size()); - CuVSMatrix.DataType dataType = expectedDataType(quantization); - int columns = expectedColumns(dimensions, quantization); - if (vectorDataset.columns() != columns) { - throw new IllegalArgumentException( - "Expected " + columns + " matrix columns, but received " + vectorDataset.columns()); - } - if (vectorDataset.dataType() != dataType) { - throw new IllegalArgumentException( - "Expected " + dataType + " matrix data, but received " + vectorDataset.dataType()); - } - return createMultiLayerHnswGraph( - size, - dimensions, - adjacencyListMatrix, - hnswLayers, - params, - selectedNodes -> createSubsetDataset(vectorDataset, selectedNodes, columns, dataType)); - } - - private static GPUBuiltHnswGraph createMultiLayerHnswGraph( - int size, - int dimensions, - CuVSMatrix adjacencyListMatrix, - int hnswLayers, - CagraIndexParams params, - SubsetDatasetFactory subsetDatasetFactory) - throws Throwable { - Objects.requireNonNull(adjacencyListMatrix, "adjacencyListMatrix"); - Objects.requireNonNull(params, "params"); - if (size < 2) { - throw new IllegalArgumentException("A multi-layer graph requires at least two vectors"); - } - if (adjacencyListMatrix.dataType() != CuVSMatrix.DataType.INT - && adjacencyListMatrix.dataType() != CuVSMatrix.DataType.UINT) { - throw new IllegalArgumentException( - "Expected INT or UINT adjacency data, but received " + adjacencyListMatrix.dataType()); - } - if (adjacencyListMatrix.size() != size) { - throw new IllegalArgumentException( - "Expected " + size + " adjacency rows, but received " + adjacencyListMatrix.size()); - } - - int degree = Math.toIntExact(adjacencyListMatrix.columns()); - if (degree <= 0) { - throw new IllegalArgumentException("The layer-0 graph must have a positive degree"); - } - int M = Math.ceilDiv(degree, 2); + int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); + // Store all layers data List layerNodes = new ArrayList<>(); List layerAdjacencies = new ArrayList<>(); // Layer 0: Use full CAGRA adjacency list - layerNodes.add(null); + layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list layerAdjacencies.add(adjacencyListMatrix); int currentLayerSize = size; int layerIndex = 1; Random random = new Random(); - Throwable failure = null; - try { - while (layerIndex < hnswLayers && currentLayerSize > 1) { - int nextLayerSize = Math.max(2, currentLayerSize / M); - SortedSet selectedNodesSet = new TreeSet<>(); + while (layerIndex < hnswLayers && currentLayerSize > 1) { + // Calculate size for next layer (1/M of current layer) + int nextLayerSize = Math.max(2, currentLayerSize / M); + // Select nodes for this layer + SortedSet selectedNodesSet = new TreeSet<>(); - if (layerIndex == 1) { - while (selectedNodesSet.size() < nextLayerSize) { - selectedNodesSet.add(random.nextInt(size)); - } - } else { - int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); - while (selectedNodesSet.size() < nextLayerSize) { - selectedNodesSet.add(prevLayerNodes[random.nextInt(prevLayerNodes.length)]); - } + if (layerIndex == 1) { + // Select from all nodes (Layer 0) + while (selectedNodesSet.size() < nextLayerSize) { + selectedNodesSet.add(random.nextInt(size)); } - - int[] selectedNodes = - selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray(); - CuVSMatrix subsetDataset = subsetDatasetFactory.create(selectedNodes); - CuVSMatrix upperAdjacency = buildCagraGraphForSubset(subsetDataset, selectedNodes, params); - try { - layerNodes.add(selectedNodes); - layerAdjacencies.add(upperAdjacency); - } catch (Throwable addFailure) { - try { - upperAdjacency.close(); - } catch (Throwable closeFailure) { - if (addFailure != closeFailure) { - addFailure.addSuppressed(closeFailure); - } - } - throw addFailure; + } else { + // Select from previous layer nodes + int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); + while (selectedNodesSet.size() < nextLayerSize) { + int idx = random.nextInt(prevLayerNodes.length); + selectedNodesSet.add(prevLayerNodes[idx]); } - - currentLayerSize = nextLayerSize; - layerIndex++; - random = new Random(new Random().nextLong()); } - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); - } catch (Throwable t) { - failure = t; - throw t; - } finally { - Throwable closeFailure = null; - for (int i = layerAdjacencies.size() - 1; i >= 1; i--) { - try { - layerAdjacencies.get(i).close(); - } catch (Throwable t) { - if (closeFailure == null) { - closeFailure = t; - } else if (closeFailure != t) { - closeFailure.addSuppressed(t); - } - } - } - if (closeFailure != null) { - if (failure != null && failure != closeFailure) { - failure.addSuppressed(closeFailure); - } else { - throw closeFailure; - } - } - } - } + // Convert to sorted array + int[] selectedNodes = + selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray(); - /** Builds a CAGRA graph for a sampled subset and remaps its ordinals to layer-0 ordinals. */ - private static CuVSMatrix buildCagraGraphForSubset( - CuVSMatrix subsetDataset, int[] selectedNodes, CagraIndexParams params) throws Throwable { - return buildCagraGraphForSubset( - subsetDataset, - selectedNodes, - params, - (dataset, indexParams) -> - CagraIndex.newBuilder(getCuVSResourcesInstance()) - .withDataset(dataset) - .withIndexParams(indexParams) - .build()); - } + layerNodes.add(selectedNodes); - static CuVSMatrix buildCagraGraphForSubset( - CuVSMatrix subsetDataset, - int[] selectedNodes, - CagraIndexParams params, - SubsetIndexFactory indexFactory) - throws Throwable { - CuVSMatrix remappedAdjacency = null; - try (Utils.OwnedIndex ownedIndex = Utils.ownDataset(subsetDataset)) { - CagraIndex subsetIndex = indexFactory.build(subsetDataset, params); - ownedIndex.transferTo(subsetIndex); - remappedAdjacency = remapSubsetGraph(subsetIndex.getGraph(), selectedNodes); - } catch (Throwable failure) { - if (remappedAdjacency != null) { - try { - remappedAdjacency.close(); - } catch (Throwable closeFailure) { - if (failure != closeFailure) { - failure.addSuppressed(closeFailure); - } + if (quantization == QuantizationType.NONE) { + // Extract vectors for selected nodes + float[][] selectedVectors = new float[nextLayerSize][]; + for (int i = 0; i < nextLayerSize; i++) { + selectedVectors[i] = (float[]) vectors.get(selectedNodes[i]); } - } - throw failure; - } - return remappedAdjacency; - } - static CuVSMatrix remapSubsetGraph(CuVSMatrix cagraGraph, int[] selectedNodes) - throws IOException { - if (cagraGraph.dataType() != CuVSMatrix.DataType.INT - && cagraGraph.dataType() != CuVSMatrix.DataType.UINT) { - throw new IOException( - "Expected INT or UINT subset adjacency data, but received " + cagraGraph.dataType()); - } - if (cagraGraph.size() != selectedNodes.length) { - throw new IOException( - "Expected " - + selectedNodes.length - + " subset adjacency rows, but received " - + cagraGraph.size()); - } - int degree = Math.toIntExact(cagraGraph.columns()); - if (degree <= 0) { - throw new IOException("The subset graph must have a positive degree"); - } - int[] remappedRow = new int[degree]; - return MatrixBuilderLifecycle.build( - CuVSMatrix.hostBuilder(selectedNodes.length, degree, CuVSMatrix.DataType.INT), - builder -> { - for (int rowOrdinal = 0; rowOrdinal < selectedNodes.length; rowOrdinal++) { - RowView row = cagraGraph.getRow(rowOrdinal); - remapSubsetAdjacencyRow(rowOrdinal, row, degree, selectedNodes, remappedRow); - builder.addVector(remappedRow); - } - return builder.build(); - }); - } + // Build CAGRA graph for this layer + layerAdjacencies.add( + buildCagraGraphForSubset( + selectedVectors, selectedNodes, 0, params, dimensions, quantization)); - static void remapSubsetAdjacencyRow( - int rowOrdinal, RowView source, int degree, int[] selectedNodes, int[] destination) - throws IOException { - if (source == null || source.size() != degree) { - throw new IOException( - "Expected " - + degree - + " neighbors for subset row " - + rowOrdinal - + ", but received " - + (source == null ? "null" : source.size())); - } - for (int column = 0; column < degree; column++) { - int subsetOrdinal = source.getAsInt(column); - if (subsetOrdinal < 0) { - destination[column] = subsetOrdinal; - } else if (subsetOrdinal >= selectedNodes.length) { - throw new IOException( - "Subset adjacency row " - + rowOrdinal - + ", column " - + column - + " contains local ordinal " - + subsetOrdinal - + " outside [0, " - + selectedNodes.length - + ")"); } else { - destination[column] = selectedNodes[subsetOrdinal]; + + // Extract vectors for selected nodes + int bytesPerVector = (dimensions + 7) / 8; + byte[][] selectedVectors = new byte[nextLayerSize][]; + for (int i = 0; i < nextLayerSize; i++) { + selectedVectors[i] = (byte[]) vectors.get(selectedNodes[i]); + } + + // Build CAGRA graph for this layer + layerAdjacencies.add( + buildCagraGraphForSubset( + selectedVectors, selectedNodes, bytesPerVector, params, dimensions, quantization)); } + + // Update for next iteration + currentLayerSize = nextLayerSize; + layerIndex++; + + // Use different seed for each layer + random = new Random(new Random().nextLong()); } + + // Create the multi-layer graph with all layers + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); } - private static CuVSMatrix createSubsetDataset( - CuVSMatrix vectors, int[] selectedNodes, int columns, CuVSMatrix.DataType dataType) - throws IOException { - return MatrixBuilderLifecycle.build( - CuVSMatrix.hostBuilder(selectedNodes.length, columns, dataType), - builder -> { - if (dataType == CuVSMatrix.DataType.FLOAT) { - float[] rowBuffer = new float[columns]; - for (int node : selectedNodes) { - RowView row = vectors.getRow(node); - validateRowWidth(row, columns, node); - row.toArray(rowBuffer); - builder.addVector(rowBuffer); - } - } else { - byte[] rowBuffer = new byte[columns]; - for (int node : selectedNodes) { - RowView row = vectors.getRow(node); - validateRowWidth(row, columns, node); - row.toArray(rowBuffer); - builder.addVector(rowBuffer); + /** + * Creates a multi-layer HNSW graph from a native matrix without copying the complete dataset to + * the Java heap. The list view copies only rows selected for an upper layer. + */ + static GPUBuiltHnswGraph createMultiLayerHnswGraph( + FieldInfo fieldInfo, + int dimensions, + CuVSMatrix adjacencyListMatrix, + CuVSMatrix vectorDataset, + int hnswLayers, + CagraIndexParams params, + QuantizationType quantization) + throws Throwable { + int size = Math.toIntExact(vectorDataset.size()); + int columns = Math.toIntExact(vectorDataset.columns()); + List vectors = + new AbstractList<>() { + @Override + public Object get(int index) { + RowView row = vectorDataset.getRow(index); + if (quantization == QuantizationType.NONE) { + float[] vector = new float[columns]; + row.toArray(vector); + return vector; } + byte[] vector = new byte[columns]; + row.toArray(vector); + return vector; } - return builder.build(); - }); - } - private static CuVSMatrix createSubsetDataset( - List vectors, int[] selectedNodes, int columns, CuVSMatrix.DataType dataType) - throws IOException { - return MatrixBuilderLifecycle.build( - CuVSMatrix.hostBuilder(selectedNodes.length, columns, dataType), - builder -> { - for (int node : selectedNodes) { - Object vector = vectors.get(node); - if (dataType == CuVSMatrix.DataType.FLOAT) { - if (!(vector instanceof float[] values) || values.length != columns) { - throw new IllegalArgumentException( - "Vector " + node + " must be a float[" + columns + "]"); - } - builder.addVector(values); - } else { - if (!(vector instanceof byte[] values) || values.length != columns) { - throw new IllegalArgumentException( - "Vector " + node + " must be a byte[" + columns + "]"); - } - builder.addVector(values); - } + @Override + public int size() { + return size; } - return builder.build(); - }); + }; + return createMultiLayerHnswGraph( + fieldInfo, + size, + dimensions, + adjacencyListMatrix, + vectors, + hnswLayers, + params, + quantization); } - private static void validateRowWidth(RowView row, int columns, int rowIndex) throws IOException { - if (row == null || row.size() != columns) { - throw new IOException( - "Expected " - + columns - + " values for vector row " - + rowIndex - + ", but received " - + (row == null ? "null" : row.size())); - } - } + /** + * Builds a CAGRA graph for a subset of binary quantized vectors + */ + private static CuVSMatrix buildCagraGraphForSubset( + Object vectors, + int[] selectedNodes, + int bytesPerVector, + CagraIndexParams params, + int dimensions, + QuantizationType quantization) + throws Throwable { - private static int expectedColumns(int dimensions, QuantizationType quantization) { - if (dimensions <= 0) { - throw new IllegalArgumentException("Vector dimensions must be positive"); + CuVSMatrix subsetDataset; + + if (quantization == QuantizationType.BINARY) { + subsetDataset = createHostByteMatrixFromArray((byte[][]) vectors, bytesPerVector); + } else if (quantization == QuantizationType.SCALAR) { + subsetDataset = createHostByteMatrixFromArray((byte[][]) vectors, dimensions); + } else { + subsetDataset = CuVSMatrix.ofArray((float[][]) vectors); } - return quantization == QuantizationType.BINARY ? Math.ceilDiv(dimensions, 8) : dimensions; - } - private static CuVSMatrix.DataType expectedDataType(QuantizationType quantization) { - Objects.requireNonNull(quantization, "quantization"); - return quantization == QuantizationType.NONE - ? CuVSMatrix.DataType.FLOAT - : CuVSMatrix.DataType.BYTE; + return buildCagraGraphForSubset(subsetDataset, selectedNodes, params); } - @FunctionalInterface - private interface SubsetDatasetFactory { - CuVSMatrix create(int[] selectedNodes) throws Throwable; - } + private static CuVSMatrix buildCagraGraphForSubset( + CuVSMatrix subsetDataset, int[] selectedNodes, CagraIndexParams params) throws Throwable { - @FunctionalInterface - interface SubsetIndexFactory { - CagraIndex build(CuVSMatrix dataset, CagraIndexParams params) throws Throwable; + // Build CAGRA index for the subset + CagraIndex subsetIndex = + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(subsetDataset) + .withIndexParams(params) + .build(); + + // Get adjacency list from subset CAGRA index + CuVSMatrix cagraGraph = subsetIndex.getGraph(); + + long numNodes = cagraGraph.size(); + long degree = cagraGraph.columns(); + + // Create a re-mapped adjacency list + int[][] remappedAdjacency = new int[(int) numNodes][(int) degree]; + + for (int i = 0; i < numNodes; i++) { + RowView rv = cagraGraph.getRow(i); + for (int j = 0; j < degree && j < rv.size(); j++) { + int subsetIndex1 = rv.getAsInt(j); + // Map subset index to original node ID + if (subsetIndex1 >= 0 && subsetIndex1 < selectedNodes.length) { + remappedAdjacency[i][j] = selectedNodes[subsetIndex1]; + } else { + // Invalid index, use self-reference + remappedAdjacency[i][j] = selectedNodes[i]; + } + } + } + + subsetIndex.close(); + return CuVSMatrix.ofArray(remappedAdjacency); } /** @@ -619,16 +464,30 @@ public static List quantizeFloatVectorsToBinary(List floatVecto } int dimensions = floatVectors.get(0).length; - BinaryQuantizer quantizer = new BinaryQuantizer(dimensions); + int numVectors = floatVectors.size(); + int bytesPerVector = (dimensions + 7) / 8; + + float[] centroids = new float[dimensions]; for (float[] vector : floatVectors) { - quantizer.add(vector); + for (int d = 0; d < dimensions; d++) { + centroids[d] += vector[d]; + } + } + for (int d = 0; d < dimensions; d++) { + centroids[d] /= numVectors; } - quantizer.finish(); - List quantizedVectors = new ArrayList<>(floatVectors.size()); + List quantizedVectors = new ArrayList<>(numVectors); for (float[] vector : floatVectors) { - byte[] quantized = new byte[Math.ceilDiv(dimensions, 8)]; - quantizer.quantize(vector, quantized); + byte[] quantized = new byte[bytesPerVector]; + for (int d = 0; d < dimensions; d++) { + boolean bit = vector[d] > centroids[d]; + int byteIndex = d / 8; + int bitIndex = d % 8; + if (bit) { + quantized[byteIndex] |= (1 << bitIndex); + } + } quantizedVectors.add(quantized); } @@ -647,159 +506,36 @@ public static List quantizeFloatVectorsToScalar(List floatVecto } int dimensions = floatVectors.get(0).length; - ScalarQuantizer quantizer = new ScalarQuantizer(dimensions); - for (float[] vector : floatVectors) { - quantizer.add(vector); - } - quantizer.finish(); - - List quantizedVectors = new ArrayList<>(floatVectors.size()); - for (float[] vector : floatVectors) { - byte[] quantized = new byte[dimensions]; - quantizer.quantize(vector, quantized); - quantizedVectors.add(quantized); - } - - return quantizedVectors; - } + int numVectors = floatVectors.size(); - static final class BinaryQuantizer { - private final float[] centroids; - private int count; - private boolean finished; + float[] minPerDim = new float[dimensions]; + float[] maxPerDim = new float[dimensions]; + Arrays.fill(minPerDim, Float.MAX_VALUE); + Arrays.fill(maxPerDim, Float.MIN_VALUE); - BinaryQuantizer(int dimensions) { - if (dimensions <= 0) { - throw new IllegalArgumentException("Vector dimensions must be positive"); - } - centroids = new float[dimensions]; - } - - void add(float[] vector) { - ensureCollecting(); - requireDimensions(vector, centroids.length); - for (int dimension = 0; dimension < vector.length; dimension++) { - centroids[dimension] += vector[dimension]; - } - count = Math.incrementExact(count); - } - - void finish() { - ensureCollecting(); - if (count == 0) { - throw new IllegalStateException("Cannot finish binary quantization without vectors"); - } - for (int dimension = 0; dimension < centroids.length; dimension++) { - centroids[dimension] /= count; - } - finished = true; - } - - void quantize(float[] vector, byte[] destination) { - ensureFinished(); - requireDimensions(vector, centroids.length); - if (destination.length != Math.ceilDiv(centroids.length, 8)) { - throw new IllegalArgumentException("Binary quantization buffer dimensions do not match"); - } - Arrays.fill(destination, (byte) 0); - for (int dimension = 0; dimension < vector.length; dimension++) { - if (vector[dimension] > centroids[dimension]) { - destination[dimension / 8] |= (byte) (1 << (dimension % 8)); - } - } - } - - int count() { - return count; - } - - private void ensureCollecting() { - if (finished) { - throw new IllegalStateException("Quantization statistics are already final"); - } - } - - private void ensureFinished() { - if (finished == false) { - throw new IllegalStateException("Quantization statistics are not final"); - } - } - } - - static final class ScalarQuantizer { - private final float[] minima; - private final float[] maxima; - private int count; - private boolean finished; - - ScalarQuantizer(int dimensions) { - if (dimensions <= 0) { - throw new IllegalArgumentException("Vector dimensions must be positive"); - } - minima = new float[dimensions]; - maxima = new float[dimensions]; - Arrays.fill(minima, Float.MAX_VALUE); - // Preserve the existing scalar quantizer's behavior for all-negative dimensions. - Arrays.fill(maxima, Float.MIN_VALUE); - } - - void add(float[] vector) { - ensureCollecting(); - requireDimensions(vector, minima.length); - for (int dimension = 0; dimension < vector.length; dimension++) { - minima[dimension] = Math.min(minima[dimension], vector[dimension]); - maxima[dimension] = Math.max(maxima[dimension], vector[dimension]); - } - count = Math.incrementExact(count); - } - - void finish() { - ensureCollecting(); - if (count == 0) { - throw new IllegalStateException("Cannot finish scalar quantization without vectors"); + for (float[] vector : floatVectors) { + for (int d = 0; d < dimensions; d++) { + minPerDim[d] = Math.min(minPerDim[d], vector[d]); + maxPerDim[d] = Math.max(maxPerDim[d], vector[d]); } - finished = true; } - void quantize(float[] vector, byte[] destination) { - ensureFinished(); - requireDimensions(vector, minima.length); - if (destination.length != minima.length) { - throw new IllegalArgumentException("Scalar quantization buffer dimensions do not match"); - } - for (int dimension = 0; dimension < vector.length; dimension++) { - float range = maxima[dimension] - minima[dimension]; + List quantizedVectors = new ArrayList<>(numVectors); + for (float[] vector : floatVectors) { + byte[] quantized = new byte[dimensions]; + for (int d = 0; d < dimensions; d++) { + float range = maxPerDim[d] - minPerDim[d]; if (range > 0) { - float normalized = (vector[dimension] - minima[dimension]) / range; + float normalized = (vector[d] - minPerDim[d]) / range; int quantizedValue = Math.round(normalized * 127.0f) - 64; - destination[dimension] = (byte) Math.max(-64, Math.min(63, quantizedValue)); + quantized[d] = (byte) Math.max(-64, Math.min(63, quantizedValue)); } else { - destination[dimension] = 0; + quantized[d] = 0; } } + quantizedVectors.add(quantized); } - int count() { - return count; - } - - private void ensureCollecting() { - if (finished) { - throw new IllegalStateException("Quantization statistics are already final"); - } - } - - private void ensureFinished() { - if (finished == false) { - throw new IllegalStateException("Quantization statistics are not final"); - } - } - } - - private static void requireDimensions(float[] vector, int dimensions) { - if (vector.length != dimensions) { - throw new IllegalArgumentException( - "Expected " + dimensions + " vector dimensions, but received " + vector.length); - } + return quantizedVectors; } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index 94e3626065..5feb211afe 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -20,9 +20,7 @@ import com.nvidia.cuvs.BruteForceIndexParams; import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraIndexParams; -import com.nvidia.cuvs.CuVSDeviceMatrix; import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.LibraryException; import java.io.IOException; import java.io.OutputStream; import java.nio.file.Files; @@ -203,18 +201,21 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro if (indexType.isCagra()) { var cagraIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); try { - CuVSDeviceMatrix cagraDataset = - createDeviceFloatMatrix(vectors, fieldInfo.getVectorDimension()); + CuVSMatrix cagraDataset = + Utils.createFloatMatrix( + vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); writeCagraIndex(cagraIndexOutputStream, cagraDataset); - } catch (RecoverableCagraConstructionException recoverable) { - ensureFallbackOutputUnchanged(cagraIndexOffset, cuvsIndex.getFilePointer(), recoverable); + } catch (Throwable t) { + // Fallback to brute force in a few cases, for now. + // Log it to make it more obvious that this is what is happening. info( infoStream, COMPONENT, "CAGRA build failed for field \"" + fieldInfo.name + "\", falling back to a brute force index: " - + recoverable.getCause()); + + t); + Utils.handleThrowableWithIgnore(t, t.getMessage()); indexType = IndexType.BRUTE_FORCE; } cagraIndexLength = cuvsIndex.getFilePointer() - cagraIndexOffset; @@ -223,7 +224,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro if (indexType.isBruteForce()) { var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); CuVSMatrix bruteforceDataset = - Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); + Utils.createFloatMatrix( + vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); writeBruteForceIndex(bruteForceIndexOutputStream, bruteforceDataset); bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; @@ -240,117 +242,38 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro } } - /** - * Creates GPU-search input directly on the device. Accelerated HNSW intentionally continues to - * use the shared host-backed matrix helpers because it does not retain a searchable device - * dataset. - */ - private CuVSDeviceMatrix createDeviceFloatMatrix(List vectors, int dimensions) - throws Throwable { - return prepareCagraDataset( - () -> - CuVSMatrix.deviceBuilder( - getCuVSResourcesInstance(), vectors.size(), dimensions, CuVSMatrix.DataType.FLOAT), - builder -> { - for (float[] vector : vectors) { - builder.addVector(vector); - } - }); - } - - @FunctionalInterface - interface CagraDatasetBuilderFactory { - CuVSMatrix.Builder create() throws Throwable; - } - - @FunctionalInterface - interface CagraDatasetPopulation { - void populate(CuVSMatrix.Builder builder) throws Throwable; - } - - /** - * Builds the device input and applies the pre-persistence fallback policy only after builder - * cleanup succeeds. A builder cleanup failure is fatal; after ownership has transferred, this - * method also closes the dataset before propagating that failure. A builder-factory failure is - * also fatal because no cleanup handle was returned. - */ - static CuVSDeviceMatrix prepareCagraDataset( - CagraDatasetBuilderFactory builderFactory, CagraDatasetPopulation population) - throws Throwable { - CuVSMatrix.Builder builder = null; - CuVSDeviceMatrix dataset = null; - Throwable operationFailure = null; - boolean factoryCompleted = false; - try { - builder = - Objects.requireNonNull(builderFactory.create(), "CAGRA dataset builder must not be null"); - factoryCompleted = true; - population.populate(builder); - dataset = builder.build(); - } catch (Throwable failure) { - operationFailure = failure; - } - - Throwable cleanupFailure = closeResource(builder, null); - if (cleanupFailure != null && dataset != null) { - cleanupFailure = closeResource(dataset, cleanupFailure); - } - - Throwable failure = combineOperationAndCleanupFailures(operationFailure, cleanupFailure); - if (failure != null) { - throw classifyCagraWriteFailure( - failure, false, cleanupFailure != null || factoryCompleted == false); - } - return dataset; - } - /** * Builds and writes the CAGRA index. * * @param os Instance of the OutputStream - * @param dataset device-backed matrix holding the dataset + * @param dataset The instance of CuVSMatrix holding the dataset * @throws Throwable */ - private void writeCagraIndex(OutputStream os, CuVSDeviceMatrix dataset) throws Throwable { - CagraIndex index = null; - AutoCloseable indexDataset = null; - Throwable failure = null; - CagraWriteContext writeContext = new CagraWriteContext(); - try { - CagraIndexParams params = - CagraIndexParamsFactory.create(gpuSearchParams, dataset.size(), dataset.columns()); - index = - CagraIndex.newBuilder(getCuVSResourcesInstance()) - .withDataset(dataset) - .withIndexParams(params) - .build(); + private void writeCagraIndex(OutputStream os, CuVSMatrix dataset) throws Throwable { + CagraIndexParams params = + CagraIndexParamsFactory.create(gpuSearchParams, dataset.size(), dataset.columns()); + try (CagraIndex index = + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(dataset) + .withIndexParams(params) + .build(); + var deviceVectors = dataset.toDevice(getCuVSResourcesInstance())) { /* * cuVS rejects makePaddedDataset for a device matrix whose rows already sit at the required - * stride, and asks for a view over that storage instead. + * stride, and asks for a view over that storage instead. Copying would be pointless there + * anyway, so pick the factory that matches the layout. */ - if (CagraIndex.isPaddedDataset(dataset)) { - var indexDatasetView = index.makePaddedDatasetView(dataset); - indexDataset = indexDatasetView; - index.updateDataset(indexDatasetView); + if (CagraIndex.isPaddedDataset(deviceVectors)) { + try (var indexDatasetView = index.makePaddedDatasetView(deviceVectors)) { + index.updateDataset(indexDatasetView); + index.serialize(os); + } } else { - var paddedDataset = index.makePaddedDataset(dataset); - indexDataset = paddedDataset; - index.updateDataset(paddedDataset); + try (var indexDataset = index.makePaddedDataset(deviceVectors)) { + index.updateDataset(indexDataset); + index.serialize(os); + } } - - // Invoking the serializer is the persistence boundary. No failure from this point can fall - // back safely, even if a particular serializer implementation happens not to write bytes. - writeContext.beginPersistence(); - index.serialize(os); - } catch (Throwable t) { - failure = t; - } - - Throwable cleanupFailure = closeCagraResources(index, dataset, indexDataset); - failure = combineOperationAndCleanupFailures(failure, cleanupFailure); - if (failure != null) { - throw classifyCagraWriteFailure( - failure, writeContext.persistenceStarted(), cleanupFailure != null); } } @@ -366,104 +289,13 @@ private void writeBruteForceIndex(OutputStream os, CuVSMatrix dataset) throws Th new BruteForceIndexParams.Builder() .withNumWriterThreads(gpuSearchParams.getWriterThreads()) .build(); - try (dataset; - var index = - BruteForceIndex.newBuilder(getCuVSResourcesInstance()) - .withIndexParams(params) - .withDataset(dataset) - .build()) { - index.serialize(os); - } - } - - private static Throwable closeResource(AutoCloseable resource, Throwable failure) { - if (resource == null) { - return failure; - } - try { - resource.close(); - } catch (Throwable closeFailure) { - return addFailure(failure, closeFailure); - } - return failure; - } - - /** Closes CAGRA resources in dependency order and returns the first cleanup failure. */ - static Throwable closeCagraResources( - AutoCloseable index, AutoCloseable originalDataset, AutoCloseable indexDataset) { - Throwable failure = null; - try { - if (index == null) { - originalDataset.close(); - } else { - Utils.closeIndexWithDatasetFallback(index, originalDataset); - } - } catch (Throwable closeFailure) { - failure = addFailure(failure, closeFailure); - } - return closeResource(indexDataset, failure); - } - - static Throwable combineOperationAndCleanupFailures( - Throwable operationFailure, Throwable cleanupFailure) { - return addFailure(operationFailure, cleanupFailure); - } - - private static Throwable addFailure(Throwable primary, Throwable secondary) { - if (secondary == null) { - return primary; - } - if (primary == null) { - return secondary; - } - if (primary != secondary) { - primary.addSuppressed(secondary); - } - return primary; - } - - /** Signals the only failure for which the caller may safely build a brute-force index instead. */ - static final class RecoverableCagraConstructionException extends Exception { - private static final long serialVersionUID = 1L; - - RecoverableCagraConstructionException(LibraryException cause) { - super(cause); - } - } - - static Throwable classifyCagraWriteFailure( - Throwable failure, boolean persistenceStarted, boolean cleanupFailed) { - if (failure instanceof LibraryException libraryFailure - && persistenceStarted == false - && cleanupFailed == false - && failure.getSuppressed().length == 0) { - return new RecoverableCagraConstructionException(libraryFailure); - } - return failure; - } - - static void ensureFallbackOutputUnchanged(long expected, long actual, Throwable failure) - throws IOException { - if (actual != expected) { - throw new IOException( - "Cannot fall back after the CAGRA output position changed from " - + expected - + " to " - + actual, - failure); - } - } - - static final class CagraWriteContext { - private boolean persistenceStarted; - - void beginPersistence() { - persistenceStarted = true; - } - - boolean persistenceStarted() { - return persistenceStarted; - } + var index = + BruteForceIndex.newBuilder(getCuVSResourcesInstance()) + .withIndexParams(params) + .withDataset(dataset) + .build(); + index.serialize(os); + index.close(); } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index fff21b62fa..7e9f888e32 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -9,8 +9,6 @@ import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.RowView; import java.util.ArrayList; -import java.util.Arrays; -import java.util.HashSet; import java.util.List; import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.NeighborArray; @@ -29,7 +27,6 @@ public class GPUBuiltHnswGraph extends HnswGraph { // Store layers data - each layer has its own nodes and adjacency lists private final List layerNodes; private final List layerNeighbors; - private final boolean[] sortedUpperLayers; // Layer 0 is special - it contains all nodes private final NeighborArray[] layer0Neighbors; @@ -44,22 +41,12 @@ public class GPUBuiltHnswGraph extends HnswGraph { */ public GPUBuiltHnswGraph( int size, int dimensions, List layerNodes, List layerAdjacencies) { - if (size < 0) { - throw new IllegalArgumentException("Graph size must not be negative"); - } - if (dimensions <= 0) { - throw new IllegalArgumentException("Vector dimensions must be positive"); - } - if (layerAdjacencies.isEmpty() || layerNodes.size() != layerAdjacencies.size()) { - throw new IllegalArgumentException( - "Layer node and adjacency lists must have the same non-zero size"); - } + this.size = size; this.dimensions = dimensions; this.numLevels = layerAdjacencies.size(); this.layerNodes = new ArrayList<>(); this.layerNeighbors = new ArrayList<>(); - this.sortedUpperLayers = new boolean[Math.max(0, numLevels - 1)]; // Process Layer 0 (base layer with all nodes) CuVSMatrix layer0Adjacency = layerAdjacencies.get(0); @@ -67,12 +54,7 @@ public GPUBuiltHnswGraph( // Process higher layers (1 to numLevels-1) for (int level = 1; level < numLevels; level++) { - int[] suppliedNodes = layerNodes.get(level); - if (suppliedNodes == null) { - throw new IllegalArgumentException("Missing node ordinals for level " + level); - } - int[] nodes = suppliedNodes.clone(); - sortedUpperLayers[level - 1] = validateLayerNodes(nodes, level); + int[] nodes = layerNodes.get(level); CuVSMatrix adjacency = layerAdjacencies.get(level); this.layerNodes.add(nodes); this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length)); @@ -87,75 +69,21 @@ public GPUBuiltHnswGraph( * @return the NeighborArray */ private NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size) { - if (adjacency.dataType() != CuVSMatrix.DataType.INT - && adjacency.dataType() != CuVSMatrix.DataType.UINT) { - throw new IllegalArgumentException( - "Expected INT or UINT adjacency data, but received " + adjacency.dataType()); - } - if (adjacency.size() != size) { - throw new IllegalArgumentException( - "Expected " + size + " adjacency rows, but received " + adjacency.size()); - } - int degree = Math.toIntExact(adjacency.columns()); - if (degree <= 0) { - throw new IllegalArgumentException("Adjacency matrices must have a positive degree"); - } NeighborArray[] neighbors = new NeighborArray[size]; for (int i = 0; i < size; i++) { RowView rv = adjacency.getRow(i); - if (rv == null || rv.size() != degree) { - throw new IllegalArgumentException( - "Expected " - + degree - + " neighbors for adjacency row " - + i - + ", but received " - + (rv == null ? "null" : rv.size())); - } - neighbors[i] = new NeighborArray(degree, true); - for (int j = 0; j < degree; j++) { - int neighbor = rv.getAsInt(j); - if (neighbor < 0) { - continue; - } - if (neighbor >= this.size) { - throw new IllegalArgumentException( - "Adjacency row " - + i - + " contains ordinal " - + neighbor - + " outside graph size " - + this.size); + if (rv != null && rv.size() > 0) { + neighbors[i] = new NeighborArray((int) rv.size(), true); + for (int j = 0; j < rv.size(); j++) { + neighbors[i].addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); } - neighbors[i].addInOrder(neighbor, 1.0f - (j * 0.001f)); + } else { + neighbors[i] = new NeighborArray(0, true); } } return neighbors; } - private boolean validateLayerNodes(int[] nodes, int level) { - boolean sorted = true; - int previous = -1; - for (int node : nodes) { - if (node < 0 || node >= size) { - throw new IllegalArgumentException( - "Level " + level + " contains node ordinal " + node + " outside [0, " + size + ")"); - } - sorted &= node > previous; - previous = node; - } - if (sorted == false) { - var uniqueNodes = new HashSet(nodes.length); - for (int node : nodes) { - if (uniqueNodes.add(node) == false) { - throw new IllegalArgumentException( - "Level " + level + " contains duplicate node ordinal " + node); - } - } - } - return sorted; - } - /** * Get all nodes on a given level as node 0th ordinals. */ @@ -183,13 +111,11 @@ public NeighborArray getNeighbors(int level, int node) { } else if (level > 0 && level < numLevels) { int[] nodes = layerNodes.get(level - 1); NeighborArray[] neighbors = layerNeighbors.get(level - 1); - if (sortedUpperLayers[level - 1]) { - int index = Arrays.binarySearch(nodes, node); - return index >= 0 ? neighbors[index] : null; - } - for (int index = 0; index < nodes.length; index++) { - if (nodes[index] == node) { - return neighbors[index]; + + // Find the index of this node in the layer + for (int i = 0; i < nodes.length; i++) { + if (nodes[i] == node) { + return neighbors[i]; } } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index bc7124c409..863877a9e3 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -142,8 +142,8 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException } /** - * Flush/sorting path: builds a host matrix from the heap vectors, then delegates - * to {@link #writeFieldInternal(FieldInfo, CuVSMatrix)}. + * Flush/sorting path: builds a host matrix from the heap vectors, then delegates to {@link + * #writeNonTrivialField(FieldInfo, CuVSMatrix)}. * * @param fieldInfo instance of FieldInfo that has the field description * @param vectors vectors to index @@ -153,26 +153,21 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro if (writeTrivialField(fieldInfo, vectors.size())) { return; } - CuVSMatrix dataset = Utils.createFloatMatrix(vectors, fieldInfo.getVectorDimension()); - writeFieldInternal(fieldInfo, dataset); + CuVSMatrix dataset = Utils.createHostFloatMatrix(vectors, fieldInfo.getVectorDimension()); + writeNonTrivialField(fieldInfo, dataset); } /** - * Builds the intermediate CAGRA index and builds and writes the HNSW index. - * Single implementation used by both the flush and merge paths. The dataset is a - * {@link CuVSMatrix} (host-backed on the merge path) so the full set of vectors is - * never double-materialised on the Java heap. + * Builds the intermediate CAGRA index and writes the HNSW index. This non-trivial path is shared + * by flushes and merges after zero- and one-vector cases have been handled. * * @param fieldInfo instance of FieldInfo that has the field description - * @param dataset matrix of all vectors to index + * @param dataset matrix of all vectors to index * @throws IOException */ - private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { - try (Utils.OwnedIndex ownedIndex = Utils.ownDataset(dataset)) { - int size = Math.toIntExact(dataset.size()); - if (writeTrivialField(fieldInfo, size)) { - return; - } + private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { + try { + int size = (int) dataset.size(); CagraIndexParams params = CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); CagraIndex cagraIndex = @@ -180,11 +175,11 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws .withDataset(dataset) .withIndexParams(params) .build(); - ownedIndex.transferTo(cagraIndex); CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); int dimensions = fieldInfo.getVectorDimension(); GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( + fieldInfo, dimensions, adjacencyListMatrix, dataset, @@ -203,8 +198,9 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws size, hnswGraph, graphLevelNodeOffsets); + cagraIndex.close(); } catch (Throwable t) { - throw Utils.handleThrowable(t); + Utils.handleThrowable(t); } } @@ -294,10 +290,9 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { } /** - * Streams merged vectors directly into a native host-memory matrix (CuVSHostMatrix) - * without materialising a List on the Java heap, then calls writeFieldInternal. - * This avoids the double-copy OOM (heap list + native matrix simultaneously) that - * occurs when force-merging large segments. + * Streams merged vectors directly into a native host-memory matrix without materializing a + * {@code List} on the Java heap. This avoids retaining both the heap list and native + * matrix during a force merge. */ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { try { @@ -305,19 +300,19 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws if (writeTrivialField(fieldInfo, size)) { return; } - int dims = fieldInfo.getVectorDimension(); FloatVectorValues mergedVectors = KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); - try (CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT)) { - KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); - for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { - builder.addVector(mergedVectors.vectorValue(it.index())); - } - writeFieldInternal(fieldInfo, builder.build()); + int dims = fieldInfo.getVectorDimension(); + CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT); + KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); + for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { + builder.addVector(mergedVectors.vectorValue(it.index())); } + CuVSHostMatrix dataset = builder.build(); + writeNonTrivialField(fieldInfo, dataset); } catch (Throwable t) { - throw Utils.handleThrowable(t); + Utils.handleThrowable(t); } } @@ -329,7 +324,7 @@ private static int countMergedVectors(FieldInfo fieldInfo, MergeState mergeState int count = 0; KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { - count = Math.incrementExact(count); + count++; } return count; } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 4bdf8a21cc..61cc31545a 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -7,6 +7,7 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createMultiLayerHnswGraph; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createSingleVectorHnswGraph; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.printInfoStream; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToBinary; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeEmpty; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeGraph; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeMeta; @@ -23,7 +24,6 @@ import com.nvidia.cuvs.CagraIndex; import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.BinaryQuantizer; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; import java.io.IOException; import java.util.ArrayList; @@ -146,45 +146,38 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException * @throws IOException */ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { - int size = vectors.size(); - if (writeTrivialField(fieldInfo, size)) { + if (vectors.size() == 0) { + writeEmpty(fieldInfo, hnswMeta); return; } + try { int dimensions = fieldInfo.getVectorDimension(); - int bytesPerVector = Math.ceilDiv(dimensions, 8); - try (CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(size, bytesPerVector, CuVSMatrix.DataType.BYTE)) { - for (byte[] vector : vectors) { - builder.addVector(vector); - } - writeFieldInternal(fieldInfo, builder.build()); - } - } catch (Throwable t) { - throw Utils.handleThrowable(t); - } - } + int bytesPerVector = (dimensions + 7) / 8; - /** Builds and writes an index from an owned binary-vector matrix. */ - private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { - try (Utils.OwnedIndex ownedIndex = Utils.ownDataset(dataset)) { - int size = Math.toIntExact(dataset.size()); - if (writeTrivialField(fieldInfo, size)) { + CuVSMatrix dataset = Utils.createHostByteMatrix(vectors, bytesPerVector); + + if (dataset.size() < 2) { + writeSingleVectorGraph(fieldInfo, vectors); return; } - int dimensions = fieldInfo.getVectorDimension(); + CagraIndexParams params = CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); + CagraIndex cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()) .withDataset(dataset) .withIndexParams(params) .build(); - ownedIndex.transferTo(cagraIndex); CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); + int size = (int) dataset.size(); + + // Create multi-layer HNSW graph from CAGRA GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( + fieldInfo, dimensions, adjacencyListMatrix, dataset, @@ -193,8 +186,11 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws QuantizationType.BINARY); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + // Write the graph to the vector index int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + + // Write metadata writeMeta( hnswVectorIndex, hnswMeta, @@ -204,22 +200,12 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws size, hnswGraph, graphLevelNodeOffsets); - } catch (Throwable t) { - throw Utils.handleThrowable(t); - } - } - /** Writes the empty or one-vector representation, if {@code size} is trivial. */ - private boolean writeTrivialField(FieldInfo fieldInfo, int size) throws IOException { - if (size == 0) { - writeEmpty(fieldInfo, hnswMeta); - return true; - } - if (size == 1) { - writeSingleVectorGraph(fieldInfo); - return true; + cagraIndex.close(); + + } catch (Throwable t) { + Utils.handleThrowable(t); } - return false; } /** @@ -273,9 +259,11 @@ private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) thr * Builds and writes a single vector graph. * * @param fieldInfo instance of FieldInfo + * @param vectors the list of binary quantized vectors * @throws IOException I/O Exceptions */ - private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { + private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) + throws IOException { // Workaround for CAGRA not supporting single vector indexes try { int size = 1; @@ -322,60 +310,20 @@ public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOE */ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { try { - int dimensions = fieldInfo.getVectorDimension(); - BinaryQuantizer quantizer = collectBinaryMergeStats(fieldInfo, mergeState, dimensions); - if (writeTrivialField(fieldInfo, quantizer.count())) { - return; - } - - int bytesPerVector = Math.ceilDiv(dimensions, 8); - FloatVectorValues mergedVectors = + FloatVectorValues mergedVectorValues = KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); - try (CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(quantizer.count(), bytesPerVector, CuVSMatrix.DataType.BYTE)) { - byte[] quantized = new byte[bytesPerVector]; - int encodedCount = 0; - KnnVectorValues.DocIndexIterator iterator = mergedVectors.iterator(); - for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) { - if (encodedCount >= quantizer.count()) { - throw new IOException( - "Merged vector count changed between quantization passes: expected " - + quantizer.count() - + ", received more rows"); - } - quantizer.quantize(mergedVectors.vectorValue(iterator.index()), quantized); - builder.addVector(quantized); - encodedCount = Math.incrementExact(encodedCount); - } - if (encodedCount != quantizer.count()) { - throw new IOException( - "Merged vector count changed between quantization passes: expected " - + quantizer.count() - + ", received " - + encodedCount); + + if (mergedVectorValues != null) { + List floatVectors = new ArrayList<>(); + KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + floatVectors.add(mergedVectorValues.vectorValue(iter.index()).clone()); } - writeFieldInternal(fieldInfo, builder.build()); + writeFieldInternal(fieldInfo, quantizeFloatVectorsToBinary(floatVectors)); } } catch (Throwable t) { - throw Utils.handleThrowable(t); - } - } - - private static BinaryQuantizer collectBinaryMergeStats( - FieldInfo fieldInfo, MergeState mergeState, int dimensions) throws IOException { - FloatVectorValues mergedVectors = - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); - BinaryQuantizer quantizer = new BinaryQuantizer(dimensions); - if (mergedVectors != null) { - KnnVectorValues.DocIndexIterator iterator = mergedVectors.iterator(); - for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) { - quantizer.add(mergedVectors.vectorValue(iterator.index())); - } - } - if (quantizer.count() > 0) { - quantizer.finish(); + Utils.handleThrowable(t); } - return quantizer; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 0ecb660663..ab76df4766 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -7,6 +7,7 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createMultiLayerHnswGraph; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.createSingleVectorHnswGraph; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.printInfoStream; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.quantizeFloatVectorsToScalar; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeEmpty; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeGraph; import static com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.writeMeta; @@ -24,7 +25,6 @@ import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; -import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.ScalarQuantizer; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -150,6 +150,14 @@ private static byte signedToUnsignedByte(byte signedByte) { return (byte) (signedByte & 0xFF); } + private static byte[] convertSignedToUnsigned(byte[] signedVector) { + byte[] unsignedVector = new byte[signedVector.length]; + for (int i = 0; i < signedVector.length; i++) { + unsignedVector[i] = signedToUnsignedByte(signedVector[i]); + } + return unsignedVector; + } + /** * Builds the intermediate CAGRA index and builds and writes the HNSW index. * @@ -158,50 +166,43 @@ private static byte signedToUnsignedByte(byte signedByte) { * @throws IOException */ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { - int size = vectors.size(); - if (writeTrivialField(fieldInfo, size)) { + if (vectors.size() == 0) { + writeEmpty(fieldInfo, hnswMeta); return; } + try { int dimensions = fieldInfo.getVectorDimension(); - try (CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(size, dimensions, CuVSMatrix.DataType.BYTE)) { - byte[] unsignedVector = new byte[dimensions]; - for (Object value : vectors) { - if (!(value instanceof byte[] signedVector) || signedVector.length != dimensions) { - throw new IllegalArgumentException( - "Expected scalar-quantized byte[" + dimensions + "] vector"); - } - copySignedToUnsigned(signedVector, unsignedVector); - builder.addVector(unsignedVector); - } - writeFieldInternal(fieldInfo, builder.build()); + + // Convert 7-bit signed bytes to 8-bit unsigned bytes for cuVS compatibility + List unsignedVectors = new ArrayList<>(vectors.size()); + for (Object signedVector : vectors) { + unsignedVectors.add(convertSignedToUnsigned((byte[]) signedVector)); } - } catch (Throwable t) { - throw Utils.handleThrowable(t); - } - } - /** Builds and writes an index from an owned scalar-vector matrix. */ - private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { - try (Utils.OwnedIndex ownedIndex = Utils.ownDataset(dataset)) { - int size = Math.toIntExact(dataset.size()); - if (writeTrivialField(fieldInfo, size)) { + // Create CuVSMatrix with BYTE data type (unsigned bytes) + CuVSMatrix dataset = Utils.createHostByteMatrix(unsignedVectors, dimensions); + + if (dataset.size() < 2) { + writeSingleVectorGraph(fieldInfo, unsignedVectors); return; } - int dimensions = fieldInfo.getVectorDimension(); + CagraIndexParams params = CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); + CagraIndex cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()) .withDataset(dataset) .withIndexParams(params) .build(); - ownedIndex.transferTo(cagraIndex); CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); + + int size = (int) dataset.size(); GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( + fieldInfo, dimensions, adjacencyListMatrix, dataset, @@ -210,8 +211,13 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws QuantizationType.SCALAR); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); + + // Write the graph to the vector index int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; + + // Write metadata writeMeta( hnswVectorIndex, hnswMeta, @@ -221,22 +227,11 @@ private void writeFieldInternal(FieldInfo fieldInfo, CuVSMatrix dataset) throws size, hnswGraph, graphLevelNodeOffsets); - } catch (Throwable t) { - throw Utils.handleThrowable(t); - } - } - /** Writes the empty or one-vector representation, if {@code size} is trivial. */ - private boolean writeTrivialField(FieldInfo fieldInfo, int size) throws IOException { - if (size == 0) { - writeEmpty(fieldInfo, hnswMeta); - return true; - } - if (size == 1) { - writeSingleVectorGraph(fieldInfo); - return true; + cagraIndex.close(); + } catch (Throwable t) { + Utils.handleThrowable(t); } - return false; } /** @@ -290,9 +285,11 @@ private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) thr * Builds and writes a single vector graph. * * @param fieldInfo instance of FieldInfo + * @param vectors the list of scalar quantized vectors (already converted to unsigned) * @throws IOException I/O Exceptions */ - private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { + private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) + throws IOException { // Workaround for CAGRA not supporting single vector indexes try { int size = 1; @@ -337,65 +334,19 @@ public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOE */ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws IOException { try { - int dimensions = fieldInfo.getVectorDimension(); - ScalarQuantizer quantizer = collectScalarMergeStats(fieldInfo, mergeState, dimensions); - if (writeTrivialField(fieldInfo, quantizer.count())) { - return; - } - - FloatVectorValues mergedVectors = + FloatVectorValues mergedVectorValues = KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); - try (CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(quantizer.count(), dimensions, CuVSMatrix.DataType.BYTE)) { - byte[] quantized = new byte[dimensions]; - int encodedCount = 0; - KnnVectorValues.DocIndexIterator iterator = mergedVectors.iterator(); - for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) { - if (encodedCount >= quantizer.count()) { - throw new IOException( - "Merged vector count changed between quantization passes: expected " - + quantizer.count() - + ", received more rows"); - } - quantizer.quantize(mergedVectors.vectorValue(iterator.index()), quantized); - copySignedToUnsigned(quantized, quantized); - builder.addVector(quantized); - encodedCount = Math.incrementExact(encodedCount); - } - if (encodedCount != quantizer.count()) { - throw new IOException( - "Merged vector count changed between quantization passes: expected " - + quantizer.count() - + ", received " - + encodedCount); + + if (mergedVectorValues != null) { + List floatVectors = new ArrayList<>(); + KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); + for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { + floatVectors.add(mergedVectorValues.vectorValue(iter.index()).clone()); } - writeFieldInternal(fieldInfo, builder.build()); + writeFieldInternal(fieldInfo, quantizeFloatVectorsToScalar(floatVectors)); } } catch (Throwable t) { - throw Utils.handleThrowable(t); - } - } - - private static ScalarQuantizer collectScalarMergeStats( - FieldInfo fieldInfo, MergeState mergeState, int dimensions) throws IOException { - FloatVectorValues mergedVectors = - KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); - ScalarQuantizer quantizer = new ScalarQuantizer(dimensions); - if (mergedVectors != null) { - KnnVectorValues.DocIndexIterator iterator = mergedVectors.iterator(); - for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) { - quantizer.add(mergedVectors.vectorValue(iterator.index())); - } - } - if (quantizer.count() > 0) { - quantizer.finish(); - } - return quantizer; - } - - private static void copySignedToUnsigned(byte[] source, byte[] destination) { - for (int i = 0; i < source.length; i++) { - destination[i] = signedToUnsignedByte(source[i]); + Utils.handleThrowable(t); } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MatrixBuilderLifecycle.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MatrixBuilderLifecycle.java deleted file mode 100644 index 4712c78ac3..0000000000 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/MatrixBuilderLifecycle.java +++ /dev/null @@ -1,64 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import com.nvidia.cuvs.CuVSMatrix; -import java.io.IOException; -import java.util.Objects; - -/** Ownership-aware matrix construction shared by the Lucene ingestion paths. */ -final class MatrixBuilderLifecycle { - - private MatrixBuilderLifecycle() {} - - @FunctionalInterface - interface BuildOperation { - T build(CuVSMatrix.Builder builder) throws Throwable; - } - - static T build(CuVSMatrix.Builder builder, BuildOperation operation) - throws IOException { - T matrix = null; - Throwable operationFailure = null; - try { - matrix = - Objects.requireNonNull(operation.build(builder), "Matrix builder must not return null"); - } catch (Throwable failure) { - operationFailure = failure; - } - - Throwable cleanupFailure = closeResource(builder, null); - if (cleanupFailure != null && matrix != null) { - cleanupFailure = closeResource(matrix, cleanupFailure); - } - Throwable failure = addFailure(operationFailure, cleanupFailure); - if (failure != null) { - throw Utils.handleThrowable(failure); - } - return matrix; - } - - private static Throwable closeResource(AutoCloseable resource, Throwable failure) { - try { - resource.close(); - } catch (Throwable closeFailure) { - return addFailure(failure, closeFailure); - } - return failure; - } - - private static Throwable addFailure(Throwable primary, Throwable secondary) { - if (secondary == null) { - return primary; - } - if (primary == null) { - return secondary; - } - if (primary != secondary) { - primary.addSuppressed(secondary); - } - return primary; - } -} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java index cdc1459e71..91bbc38e2c 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -51,141 +51,125 @@ static RuntimeException handleThrowable(Throwable t) throws IOException { } /** - * Builds a host-memory CuVSMatrix from a list of float vectors. + * A method to build a CuVSMatrix from a list of float vectors. * - *

Copies vectors directly into a native host matrix via {@link CuVSMatrix#hostBuilder}, - * without creating an intermediate {@code float[][]} on the heap. + * Uses CuVSMatrix.Builder to copy vectors directly to device memory + * without creating intermediate heap arrays. * * @param data The float vectors - * @param dimensions The number of float elements in each vector - * @return a host-memory CuVSMatrix + * @param dimensions The number float elements in each vector + * @param resources The CuVS resources for device matrix creation + * @return an instance of CuVSMatrix */ - static CuVSMatrix createFloatMatrix(List data, int dimensions) throws IOException { - return MatrixBuilderLifecycle.build( - CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT), - builder -> { - for (float[] vector : data) { - builder.addVector(vector); - } - return builder.build(); - }); + static CuVSMatrix createFloatMatrix(List data, int dimensions, CuVSResources resources) { + // Use Builder pattern to avoid intermediate float[][] allocation + // and copy directly from List to device memory + CuVSMatrix.Builder builder = + CuVSMatrix.deviceBuilder( + resources, + data.size(), // rows (number of vectors) + dimensions, // columns (vector dimension) + CuVSMatrix.DataType.FLOAT); + + // Add vectors one by one - builder copies directly to device memory + for (float[] vector : data) { + builder.addVector(vector); + } + + return builder.build(); } /** - * Builds a host-memory CuVSMatrix from a list of byte vectors (e.g. quantized vectors). + * Builds a host-memory CuVSMatrix from a list of float vectors. * - * @param data The byte vectors (packed bits for binary quantization) - * @param bytesPerVector The number of bytes in each vector - * @return a host-memory CuVSMatrix with BYTE data type + *

Copies vectors directly into native host memory without creating an intermediate {@code + * float[][]} on the heap. + * + * @param data The float vectors + * @param dimensions The number of float elements in each vector + * @return a host-memory CuVSMatrix */ - static CuVSMatrix createByteMatrix(List data, int bytesPerVector) throws IOException { - return MatrixBuilderLifecycle.build( - CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE), - builder -> { - for (byte[] vector : data) { - builder.addVector(vector); - } - return builder.build(); - }); + static CuVSMatrix createHostFloatMatrix(List data, int dimensions) { + CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT); + for (float[] vector : data) { + builder.addVector(vector); + } + return builder.build(); } /** - * Builds a host-memory CuVSMatrix from a 2D byte array (e.g. quantized vectors). + * A method to build a CuVSMatrix from a list of byte vectors (for binary quantized vectors). * - * @param data The 2D byte array (packed bits for binary quantization) + * Uses CuVSMatrix.Builder to copy vectors directly to device memory + * without creating intermediate heap arrays. + * + * @param data The byte vectors (packed bits for binary quantization) * @param bytesPerVector The number of bytes in each vector - * @return a host-memory CuVSMatrix with BYTE data type - */ - static CuVSMatrix createByteMatrixFromArray(byte[][] data, int bytesPerVector) - throws IOException { - return MatrixBuilderLifecycle.build( - CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE), - builder -> { - for (byte[] vector : data) { - builder.addVector(vector); - } - return builder.build(); - }); - } - - /** - * Closes an index that owns {@code dataset}. If index cleanup fails before releasing the - * dataset, a direct dataset close is attempted and attached to the index failure when needed. + * @param resources The CuVS resources for device matrix creation + * @return an instance of CuVSMatrix with BYTE data type */ - static void closeIndexWithDatasetFallback(AutoCloseable index, AutoCloseable dataset) - throws Exception { - try { - index.close(); - } catch (Throwable indexCloseFailure) { - try { - dataset.close(); - } catch (Throwable datasetCloseFailure) { - if (indexCloseFailure != datasetCloseFailure) { - indexCloseFailure.addSuppressed(datasetCloseFailure); - } - } - rethrowCloseFailure(indexCloseFailure); + static CuVSMatrix createByteMatrix( + List data, int bytesPerVector, CuVSResources resources) { + // Use Builder pattern to avoid intermediate byte[][] allocation + // and copy directly from List to device memory + CuVSMatrix.Builder builder = + CuVSMatrix.deviceBuilder( + resources, + data.size(), // rows (number of vectors) + bytesPerVector, // columns (bytes per vector) + CuVSMatrix.DataType.BYTE); + + // Add vectors one by one - builder copies directly to device memory + for (byte[] vector : data) { + builder.addVector(vector); } + + return builder.build(); } - /** Starts an ownership scope for a dataset that may later be transferred to an index. */ - static OwnedIndex ownDataset(AutoCloseable dataset) { - return new OwnedIndex<>(dataset); + /** Builds a host-memory CuVSMatrix from a list of byte vectors. */ + static CuVSMatrix createHostByteMatrix(List data, int bytesPerVector) { + CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE); + for (byte[] vector : data) { + builder.addVector(vector); + } + return builder.build(); } /** - * Owns a dataset immediately and, after {@link #transferTo}, closes the owning index with a - * direct dataset-close fallback. + * A method to build a CuVSMatrix from a 2D byte array (for binary quantized vectors). + * + * @param data The 2D byte array (packed bits for binary quantization) + * @param bytesPerVector The number of bytes in each vector + * @param resources The CuVS resources for device matrix creation + * @return an instance of CuVSMatrix with BYTE data type */ - static final class OwnedIndex implements AutoCloseable { - private AutoCloseable dataset; - private I index; - private boolean closed; - - private OwnedIndex(AutoCloseable dataset) { - this.dataset = java.util.Objects.requireNonNull(dataset, "dataset"); - } - - void transferTo(I index) { - if (closed || this.index != null) { - throw new IllegalStateException("Dataset ownership has already been transferred"); - } - this.index = java.util.Objects.requireNonNull(index, "index"); - } - - I index() { - if (index == null) { - throw new IllegalStateException("Dataset ownership has not been transferred to an index"); - } - return index; - } - - @Override - public void close() throws Exception { - if (closed) { - return; - } - closed = true; - AutoCloseable ownedDataset = dataset; - I ownedIndex = index; - dataset = null; - index = null; - if (ownedIndex == null) { - ownedDataset.close(); - } else { - closeIndexWithDatasetFallback(ownedIndex, ownedDataset); - } + static CuVSMatrix createByteMatrixFromArray( + byte[][] data, int bytesPerVector, CuVSResources resources) { + CuVSMatrix.Builder builder = + CuVSMatrix.deviceBuilder( + resources, + data.length, // rows (number of vectors) + bytesPerVector, // columns (bytes per vector) + CuVSMatrix.DataType.BYTE); + + // Add vectors one by one - builder copies directly to device memory + for (byte[] vector : data) { + builder.addVector(vector); } + return builder.build(); } - private static void rethrowCloseFailure(Throwable failure) throws Exception { - if (failure instanceof Exception exception) { - throw exception; - } - if (failure instanceof Error error) { - throw error; + /** Builds a host-memory CuVSMatrix from a 2D byte array. */ + static CuVSMatrix createHostByteMatrixFromArray(byte[][] data, int bytesPerVector) { + CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE); + for (byte[] vector : data) { + builder.addVector(vector); } - throw new AssertionError("Unexpected throwable from AutoCloseable.close()", failure); + return builder.build(); } /** diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java index 9cc5ec87ed..28fa4d974e 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java @@ -455,12 +455,6 @@ private void assertTrivialLiveVectorMerge(int liveVectors) throws Exception { assertEquals(liveVectors, graph.size()); assertEquals(liveVectors == 0 ? 0 : 1, graph.numLevels()); assertEquals(liveVectors, graph.getNodesOnLevel(0).size()); - assertAllGraphOrdinalsInBounds(graph, liveVectors); - if (liveVectors == 1) { - graph.seek(0, 0); - assertEquals( - "a single-node graph must have no edges", NO_MORE_DOCS, graph.nextNeighbor()); - } ((CodecReader) leaf).getVectorReader().checkIntegrity(); IndexSearcher searcher = new IndexSearcher(reader); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWQuantizers.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWQuantizers.java deleted file mode 100644 index 52adb537c3..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWQuantizers.java +++ /dev/null @@ -1,55 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.BinaryQuantizer; -import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.ScalarQuantizer; -import java.util.List; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.junit.Test; - -public class TestAcceleratedHNSWQuantizers extends LuceneTestCase { - - @Test - public void testBinaryGoldenBytesAndReusableBufferClearing() { - float[] low = new float[10]; - float[] high = new float[] {2, 0, 2, 0, 2, 0, 2, 0, 2, 0}; - BinaryQuantizer quantizer = new BinaryQuantizer(10); - quantizer.add(low); - quantizer.add(high); - quantizer.finish(); - - byte[] scratch = new byte[] {(byte) 0xff, (byte) 0xff}; - quantizer.quantize(high, scratch); - assertArrayEquals(new byte[] {0x55, 0x01}, scratch); - - quantizer.quantize(low, scratch); - assertArrayEquals(new byte[] {0x00, 0x00}, scratch); - List quantized = AcceleratedHNSWUtils.quantizeFloatVectorsToBinary(List.of(low, high)); - assertArrayEquals(new byte[] {0x00, 0x00}, quantized.get(0)); - assertArrayEquals(new byte[] {0x55, 0x01}, quantized.get(1)); - } - - @Test - public void testScalarGoldenBytesPreserveExistingNegativeAndConstantBehavior() { - float[] first = new float[] {-4, -2, 7}; - float[] second = new float[] {-2, -2, 7}; - ScalarQuantizer quantizer = new ScalarQuantizer(3); - quantizer.add(first); - quantizer.add(second); - quantizer.finish(); - - byte[] scratch = new byte[3]; - quantizer.quantize(first, scratch); - assertArrayEquals(new byte[] {-64, -64, 0}, scratch); - quantizer.quantize(second, scratch); - assertArrayEquals(new byte[] {0, -64, 0}, scratch); - - List quantized = - AcceleratedHNSWUtils.quantizeFloatVectorsToScalar(List.of(first, second)); - assertArrayEquals(new byte[] {-64, -64, 0}, quantized.get(0)); - assertArrayEquals(new byte[] {0, -64, 0}, quantized.get(1)); - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUpperLayers.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUpperLayers.java index aa20071797..0c7b5db007 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUpperLayers.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUpperLayers.java @@ -5,21 +5,13 @@ package com.nvidia.cuvs.lucene; import com.nvidia.cuvs.CagraIndexParams; -import com.nvidia.cuvs.CuVSDeviceMatrix; -import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.RowView; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; -import java.io.IOException; import java.lang.reflect.Method; import java.lang.reflect.Modifier; -import java.util.ArrayList; -import java.util.Arrays; import java.util.List; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; import org.junit.Test; public class TestAcceleratedHNSWUpperLayers extends LuceneTestCase { @@ -43,6 +35,7 @@ public void testLegacyListOverloadDescriptorIsPresent() throws Exception { Method matrixOverload = AcceleratedHNSWUtils.class.getDeclaredMethod( "createMultiLayerHnswGraph", + FieldInfo.class, int.class, CuVSMatrix.class, CuVSMatrix.class, @@ -51,223 +44,4 @@ public void testLegacyListOverloadDescriptorIsPresent() throws Exception { QuantizationType.class); assertFalse(Modifier.isPublic(matrixOverload.getModifiers())); } - - @Test - public void testRemapPreservesSentinelAndMapsAbsoluteOrdinals() throws IOException { - int[] destination = new int[3]; - - AcceleratedHNSWUtils.remapSubsetAdjacencyRow( - 0, new IntRow(-1, 2, 0), 3, new int[] {2, 7, 11}, destination); - - assertArrayEquals(new int[] {-1, 11, 2}, destination); - } - - @Test - public void testRemapRejectsPositiveOrdinalOutsideSubset() { - IOException thrown = - assertThrows( - IOException.class, - () -> - AcceleratedHNSWUtils.remapSubsetAdjacencyRow( - 4, new IntRow(0, 3), 2, new int[] {2, 7, 11}, new int[2])); - - assertTrue(thrown.getMessage().contains("row 4, column 1")); - assertTrue(thrown.getMessage().contains("outside [0, 3)")); - } - - @Test - public void testGraphMaterializationSkipsNativeSentinels() { - IntMatrix adjacency = new IntMatrix(new int[][] {{-1, 1, 2}, {0, -1, 2}, {0, 1, -1}}); - List layerNodes = new ArrayList<>(); - layerNodes.add(null); - - GPUBuiltHnswGraph graph = new GPUBuiltHnswGraph(3, 2, layerNodes, List.of(adjacency)); - - assertArrayEquals( - new int[] {1, 2}, - Arrays.copyOf(graph.getNeighbors(0, 0).nodes(), graph.getNeighbors(0, 0).size())); - assertEquals(2, graph.getNeighbors(0, 0).size()); - } - - @Test - public void testGraphMaterializationRejectsPositiveOrdinalOutsideGraph() { - IntMatrix adjacency = new IntMatrix(new int[][] {{1}, {2}}); - List layerNodes = new ArrayList<>(); - layerNodes.add(null); - - IllegalArgumentException thrown = - assertThrows( - IllegalArgumentException.class, - () -> new GPUBuiltHnswGraph(2, 2, layerNodes, List.of(adjacency))); - - assertTrue(thrown.getMessage().contains("ordinal 2 outside graph size 2")); - } - - @Test - public void testUpperLayerNeighborLookupUsesSortedNodeOrdinals() { - IntMatrix baseAdjacency = new IntMatrix(new int[][] {{-1}, {-1}, {-1}, {-1}, {-1}, {-1}, {-1}}); - IntMatrix upperAdjacency = new IntMatrix(new int[][] {{3}, {6}, {1}}); - int[] upperNodes = new int[] {1, 3, 6}; - GPUBuiltHnswGraph graph = - new GPUBuiltHnswGraph( - 7, 2, Arrays.asList((int[]) null, upperNodes), List.of(baseAdjacency, upperAdjacency)); - - Arrays.fill(upperNodes, 0); - assertArrayEquals(new int[] {1, 3, 6}, NodesIterator.getSortedNodes(graph.getNodesOnLevel(1))); - assertArrayEquals(new int[] {3}, neighbors(graph, 1, 1)); - assertArrayEquals(new int[] {6}, neighbors(graph, 1, 3)); - assertArrayEquals(new int[] {1}, neighbors(graph, 1, 6)); - assertNull(graph.getNeighbors(1, 0)); - assertNull(graph.getNeighbors(1, 4)); - assertNull(graph.getNeighbors(1, 7)); - } - - @Test - public void testUpperLayerNeighborLookupPreservesUnsortedPublicInput() { - IntMatrix baseAdjacency = new IntMatrix(new int[][] {{-1}, {-1}, {-1}, {-1}, {-1}, {-1}, {-1}}); - IntMatrix upperAdjacency = new IntMatrix(new int[][] {{1}, {3}, {6}}); - GPUBuiltHnswGraph graph = - new GPUBuiltHnswGraph( - 7, - 2, - Arrays.asList((int[]) null, new int[] {6, 1, 3}), - List.of(baseAdjacency, upperAdjacency)); - - assertArrayEquals(new int[] {1}, neighbors(graph, 1, 6)); - assertArrayEquals(new int[] {3}, neighbors(graph, 1, 1)); - assertArrayEquals(new int[] {6}, neighbors(graph, 1, 3)); - assertNull(graph.getNeighbors(1, 2)); - } - - @Test - public void testUpperLayerNodeValidationRejectsDuplicatesAndOutOfRangeOrdinals() { - IntMatrix baseAdjacency = new IntMatrix(new int[][] {{-1}, {-1}, {-1}, {-1}}); - IntMatrix threeRows = new IntMatrix(new int[][] {{1}, {2}, {3}}); - IllegalArgumentException duplicate = - assertThrows( - IllegalArgumentException.class, - () -> - new GPUBuiltHnswGraph( - 4, - 2, - Arrays.asList((int[]) null, new int[] {1, 2, 1}), - List.of(baseAdjacency, threeRows))); - assertTrue(duplicate.getMessage().contains("duplicate node ordinal 1")); - - IntMatrix twoRows = new IntMatrix(new int[][] {{1}, {2}}); - IllegalArgumentException outOfRange = - assertThrows( - IllegalArgumentException.class, - () -> - new GPUBuiltHnswGraph( - 4, - 2, - Arrays.asList((int[]) null, new int[] {1, 4}), - List.of(baseAdjacency, twoRows))); - assertTrue(outOfRange.getMessage().contains("ordinal 4 outside [0, 4)")); - } - - private static int[] neighbors(GPUBuiltHnswGraph graph, int level, int node) { - var neighbors = graph.getNeighbors(level, node); - return Arrays.copyOf(neighbors.nodes(), neighbors.size()); - } - - private record IntRow(int... values) implements RowView { - @Override - public long size() { - return values.length; - } - - @Override - public int getAsInt(long index) { - return values[Math.toIntExact(index)]; - } - - @Override - public float getAsFloat(long index) { - throw new UnsupportedOperationException(); - } - - @Override - public byte getAsByte(long index) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(int[] array) { - System.arraycopy(values, 0, array, 0, values.length); - } - - @Override - public void toArray(float[] array) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(byte[] array) { - throw new UnsupportedOperationException(); - } - } - - private record IntMatrix(int[][] values) implements CuVSMatrix { - @Override - public long size() { - return values.length; - } - - @Override - public long columns() { - return values[0].length; - } - - @Override - public DataType dataType() { - return DataType.INT; - } - - @Override - public RowView getRow(long row) { - return new IntRow(values[Math.toIntExact(row)]); - } - - @Override - public void toArray(int[][] array) { - for (int row = 0; row < values.length; row++) { - System.arraycopy(values[row], 0, array[row], 0, values[row].length); - } - } - - @Override - public void toArray(float[][] array) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(byte[][] array) { - throw new UnsupportedOperationException(); - } - - @Override - public void toHost(CuVSHostMatrix hostMatrix) { - throw new UnsupportedOperationException(); - } - - @Override - public CuVSHostMatrix toHost() { - throw new UnsupportedOperationException(); - } - - @Override - public void toDevice(CuVSDeviceMatrix deviceMatrix, CuVSResources resources) { - throw new UnsupportedOperationException(); - } - - @Override - public CuVSDeviceMatrix toDevice(CuVSResources resources) { - throw new UnsupportedOperationException(); - } - - @Override - public void close() {} - } } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java index 5effa973b1..89d58aa682 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexAtAlignedDimensions.java @@ -11,15 +11,10 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; -import org.apache.lucene.codecs.KnnVectorsReader; -import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; import org.apache.lucene.document.Document; import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.index.CodecReader; -import org.apache.lucene.index.DirectoryReader; import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.util.LuceneTestCase; @@ -43,14 +38,14 @@ public class TestCagraIndexAtAlignedDimensions extends LuceneTestCase { @Test public void testCagraIsBuiltAtAnAlignedDimension() throws IOException { - // Deep1B's 96 floats occupy 384 bytes, an exact multiple of CAGRA's 16-byte row alignment. - assertCagraIsBuilt(96); + // 128 floats is 512 bytes, an exact multiple of the 16 byte CAGRA row alignment. + assertCagraIsBuilt(128); } @Test public void testCagraIsBuiltAtAnUnalignedDimension() throws IOException { - // 95 floats is not aligned, so the writer has to create an owning padded copy. - assertCagraIsBuilt(95); + // 127 floats is not, so the writer has to fall back to an owning padded copy. + assertCagraIsBuilt(127); } /** Indexes a segment of the given dimension and fails if the CAGRA build did not survive it. */ @@ -59,11 +54,6 @@ private void assertCagraIsBuilt(int dimension) throws IOException { RecordingInfoStream infoStream = new RecordingInfoStream(); try (Directory directory = newDirectory()) { - float[] queryVector = new float[dimension]; - for (int d = 0; d < dimension; d++) { - queryVector[d] = random().nextFloat(); - } - IndexWriterConfig config = new IndexWriterConfig() .setCodec( @@ -74,11 +64,9 @@ private void assertCagraIsBuilt(int dimension) throws IOException { try (IndexWriter writer = new IndexWriter(directory, config)) { for (int i = 0; i < 64; i++) { - float[] vector = i == 0 ? queryVector : new float[dimension]; - if (i != 0) { - for (int d = 0; d < dimension; d++) { - vector[d] = random().nextFloat(); - } + float[] vector = new float[dimension]; + for (int d = 0; d < dimension; d++) { + vector[d] = random().nextFloat(); } Document doc = new Document(); doc.add(new KnnFloatVectorField("vector", vector, VectorSimilarityFunction.EUCLIDEAN)); @@ -86,24 +74,6 @@ private void assertCagraIsBuilt(int dimension) throws IOException { } writer.commit(); } - - try (DirectoryReader reader = DirectoryReader.open(directory)) { - LeafReader leaf = getOnlyLeafReader(reader); - CuVS2510GPUVectorsReader gpuReader = gpuReader(leaf, "vector"); - CuVS2510GPUVectorsReader.FieldEntry fieldEntry = gpuReader.getFieldEntry("vector"); - assertNotNull(fieldEntry); - assertTrue("Expected a serialized CAGRA payload", fieldEntry.cagraIndexLength() > 0); - assertEquals("Did not expect a brute-force payload", 0, fieldEntry.bruteForceIndexLength()); - assertNotNull( - "Expected the CAGRA index to be loaded", gpuReader.getCagraIndexForField("vector")); - - var searcher = newSearcher(reader); - var query = new GPUKnnFloatVectorQuery("vector", queryVector, 1, null, 1, 1); - var topDocs = searcher.search(query, 1); - assertEquals(1, topDocs.scoreDocs.length); - assertTrue(topDocs.scoreDocs[0].doc >= 0); - assertTrue(topDocs.scoreDocs[0].doc < reader.maxDoc()); - } } assertTrue( @@ -114,14 +84,6 @@ private void assertCagraIsBuilt(int dimension) throws IOException { infoStream.cagraBuildFailures().isEmpty()); } - private static CuVS2510GPUVectorsReader gpuReader(LeafReader leaf, String field) { - KnnVectorsReader reader = ((CodecReader) leaf).getVectorReader(); - if (reader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { - reader = fieldsReader.getFieldReader(field); - } - return (CuVS2510GPUVectorsReader) reader; - } - /** An InfoStream that keeps the messages, so that a test can tell which index type was built. */ private static class RecordingInfoStream extends InfoStream { diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsWriterFailureHandling.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsWriterFailureHandling.java deleted file mode 100644 index 605ca8e64f..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVS2510GPUVectorsWriterFailureHandling.java +++ /dev/null @@ -1,337 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import com.nvidia.cuvs.CuVSDeviceMatrix; -import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.LibraryException; -import java.io.IOException; -import java.lang.reflect.Proxy; -import java.util.ArrayList; -import java.util.List; -import java.util.concurrent.atomic.AtomicInteger; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.junit.Test; - -public class TestCuVS2510GPUVectorsWriterFailureHandling extends LuceneTestCase { - - @Test - public void testDevicePreparationNativeFailureIsRecoverable() { - LibraryException nativeFailure = new LibraryException("device upload"); - FakeDeviceBuilder builder = new FakeDeviceBuilder(null, null); - - Throwable thrown = - assertThrows( - Throwable.class, - () -> - CuVS2510GPUVectorsWriter.prepareCagraDataset( - () -> builder, - ignored -> { - throw nativeFailure; - })); - - assertTrue(thrown instanceof CuVS2510GPUVectorsWriter.RecoverableCagraConstructionException); - assertSame(nativeFailure, thrown.getCause()); - assertEquals(0, nativeFailure.getSuppressed().length); - assertEquals(1, builder.closeCalls.get()); - } - - @Test - public void testDevicePreparationOperationAndCleanupFailureIsFatal() { - LibraryException nativeFailure = new LibraryException("device upload"); - IllegalStateException cleanupFailure = new IllegalStateException("builder close"); - FakeDeviceBuilder builder = new FakeDeviceBuilder(null, cleanupFailure); - - Throwable thrown = - assertThrows( - Throwable.class, - () -> - CuVS2510GPUVectorsWriter.prepareCagraDataset( - () -> builder, - ignored -> { - throw nativeFailure; - })); - - assertSame(nativeFailure, thrown); - assertArrayEquals(new Throwable[] {cleanupFailure}, thrown.getSuppressed()); - assertEquals(1, builder.closeCalls.get()); - } - - @Test - public void testDevicePreparationNativeFactoryFailureIsFatal() { - LibraryException factoryFailure = new LibraryException("builder factory"); - FakeDeviceBuilder builder = new FakeDeviceBuilder(null, null); - - Throwable thrown = - assertThrows( - Throwable.class, - () -> - CuVS2510GPUVectorsWriter.prepareCagraDataset( - () -> { - throw factoryFailure; - }, - ignored -> fail("population must not run"))); - - assertSame(factoryFailure, thrown); - assertFalse(thrown instanceof CuVS2510GPUVectorsWriter.RecoverableCagraConstructionException); - assertEquals(0, builder.closeCalls.get()); - } - - @Test - public void testDevicePreparationTransfersOwnershipAfterBuilderCleanup() throws Throwable { - AtomicInteger datasetCloses = new AtomicInteger(); - CuVSDeviceMatrix dataset = fakeDeviceMatrix(datasetCloses, null); - FakeDeviceBuilder builder = new FakeDeviceBuilder(dataset, null); - - CuVSDeviceMatrix prepared = - CuVS2510GPUVectorsWriter.prepareCagraDataset(() -> builder, ignored -> {}); - - assertSame(dataset, prepared); - assertEquals(1, builder.closeCalls.get()); - assertEquals(0, datasetCloses.get()); - } - - @Test - public void testDevicePreparationBuilderCleanupFailureClosesTransferredDataset() { - LibraryException cleanupFailure = new LibraryException("builder close"); - AtomicInteger datasetCloses = new AtomicInteger(); - CuVSDeviceMatrix dataset = fakeDeviceMatrix(datasetCloses, null); - FakeDeviceBuilder builder = new FakeDeviceBuilder(dataset, cleanupFailure); - - Throwable thrown = - assertThrows( - Throwable.class, - () -> CuVS2510GPUVectorsWriter.prepareCagraDataset(() -> builder, ignored -> {})); - - assertSame(cleanupFailure, thrown); - assertEquals(1, builder.closeCalls.get()); - assertEquals(1, datasetCloses.get()); - } - - @Test - public void testDevicePreparationDatasetCleanupFailureIsSuppressed() { - LibraryException builderFailure = new LibraryException("builder close"); - IllegalStateException datasetFailure = new IllegalStateException("dataset close"); - AtomicInteger datasetCloses = new AtomicInteger(); - CuVSDeviceMatrix dataset = fakeDeviceMatrix(datasetCloses, datasetFailure); - FakeDeviceBuilder builder = new FakeDeviceBuilder(dataset, builderFailure); - - Throwable thrown = - assertThrows( - Throwable.class, - () -> CuVS2510GPUVectorsWriter.prepareCagraDataset(() -> builder, ignored -> {})); - - assertSame(builderFailure, thrown); - assertArrayEquals(new Throwable[] {datasetFailure}, thrown.getSuppressed()); - assertEquals(1, builder.closeCalls.get()); - assertEquals(1, datasetCloses.get()); - } - - @Test - public void testDevicePreparationPreservesNonNativeFailures() { - IOException checkedFailure = new IOException("checked operation"); - AssertionError error = new AssertionError("operation error"); - - for (Throwable expected : List.of(checkedFailure, error)) { - FakeDeviceBuilder builder = new FakeDeviceBuilder(null, null); - Throwable thrown = - assertThrows( - Throwable.class, - () -> - CuVS2510GPUVectorsWriter.prepareCagraDataset( - () -> builder, - ignored -> { - throw expected; - })); - - assertSame(expected, thrown); - assertEquals(1, builder.closeCalls.get()); - } - } - - @Test - public void testNativeConstructionFailureIsRecoverableBeforePersistence() { - LibraryException nativeFailure = new LibraryException("native construction"); - - Throwable classified = - CuVS2510GPUVectorsWriter.classifyCagraWriteFailure(nativeFailure, false, false); - - assertTrue( - classified instanceof CuVS2510GPUVectorsWriter.RecoverableCagraConstructionException); - assertSame(nativeFailure, classified.getCause()); - } - - @Test - public void testArbitraryConstructionFailureIsNotRecoverable() { - RuntimeException programmingFailure = new IllegalStateException("programming failure"); - - Throwable classified = - CuVS2510GPUVectorsWriter.classifyCagraWriteFailure(programmingFailure, false, false); - - assertSame(programmingFailure, classified); - } - - @Test - public void testNativeFailureIsNotRecoverableAfterPersistenceStarts() { - LibraryException nativeFailure = new LibraryException("serialization"); - - Throwable classified = - CuVS2510GPUVectorsWriter.classifyCagraWriteFailure(nativeFailure, true, false); - - assertSame(nativeFailure, classified); - } - - @Test - public void testCleanupFailureDisablesFallback() { - LibraryException nativeFailure = new LibraryException("native construction"); - nativeFailure.addSuppressed(new IOException("cleanup")); - - Throwable classified = - CuVS2510GPUVectorsWriter.classifyCagraWriteFailure(nativeFailure, false, true); - - assertSame(nativeFailure, classified); - } - - @Test - public void testSuppressedFailureDisablesFallback() { - LibraryException nativeFailure = new LibraryException("native construction"); - nativeFailure.addSuppressed(new IOException("cleanup")); - - Throwable classified = - CuVS2510GPUVectorsWriter.classifyCagraWriteFailure(nativeFailure, false, false); - - assertSame(nativeFailure, classified); - } - - @Test - public void testPersistenceBoundaryIsMonotonic() { - CuVS2510GPUVectorsWriter.CagraWriteContext context = - new CuVS2510GPUVectorsWriter.CagraWriteContext(); - - assertFalse(context.persistenceStarted()); - context.beginPersistence(); - assertTrue(context.persistenceStarted()); - } - - @Test - public void testChangedOutputPositionRejectsFallback() { - Throwable failure = new LibraryException("native construction"); - - IOException thrown = - assertThrows( - IOException.class, - () -> CuVS2510GPUVectorsWriter.ensureFallbackOutputUnchanged(10L, 11L, failure)); - - assertSame(failure, thrown.getCause()); - assertTrue(thrown.getMessage().contains("changed from 10 to 11")); - } - - @Test - public void testUnchangedOutputPositionAllowsFallback() throws IOException { - CuVS2510GPUVectorsWriter.ensureFallbackOutputUnchanged( - 10L, 10L, new LibraryException("native construction")); - } - - @Test - public void testCagraResourcesCloseInDependencyOrder() { - List closeOrder = new ArrayList<>(); - AtomicInteger directDatasetCloses = new AtomicInteger(); - - Throwable failure = - CuVS2510GPUVectorsWriter.closeCagraResources( - () -> closeOrder.add("index"), - directDatasetCloses::incrementAndGet, - () -> closeOrder.add("padded-dataset")); - - assertNull(failure); - assertEquals(List.of("index", "padded-dataset"), closeOrder); - assertEquals(0, directDatasetCloses.get()); - } - - @Test - public void testBodyAndCleanupFailuresPreserveOrderAndSuppression() { - List closeOrder = new ArrayList<>(); - IOException bodyFailure = new IOException("serialize"); - RuntimeException indexFailure = new RuntimeException("index close"); - RuntimeException datasetFailure = new RuntimeException("dataset close"); - RuntimeException paddedFailure = new RuntimeException("padded close"); - - Throwable cleanupFailure = - CuVS2510GPUVectorsWriter.closeCagraResources( - failingCloseable("index", closeOrder, indexFailure), - failingCloseable("dataset", closeOrder, datasetFailure), - failingCloseable("padded-dataset", closeOrder, paddedFailure)); - Throwable combined = - CuVS2510GPUVectorsWriter.combineOperationAndCleanupFailures(bodyFailure, cleanupFailure); - - assertSame(bodyFailure, combined); - assertEquals(List.of("index", "dataset", "padded-dataset"), closeOrder); - assertArrayEquals(new Throwable[] {indexFailure}, bodyFailure.getSuppressed()); - assertArrayEquals( - new Throwable[] {datasetFailure, paddedFailure}, indexFailure.getSuppressed()); - } - - private static AutoCloseable failingCloseable( - String name, List closeOrder, RuntimeException failure) { - return () -> { - closeOrder.add(name); - throw failure; - }; - } - - private static CuVSDeviceMatrix fakeDeviceMatrix( - AtomicInteger closeCalls, RuntimeException closeFailure) { - return (CuVSDeviceMatrix) - Proxy.newProxyInstance( - CuVSDeviceMatrix.class.getClassLoader(), - new Class[] {CuVSDeviceMatrix.class}, - (proxy, method, args) -> { - if (method.getName().equals("close") && method.getParameterCount() == 0) { - closeCalls.incrementAndGet(); - if (closeFailure != null) { - throw closeFailure; - } - return null; - } - throw new AssertionError("Unexpected device-matrix call: " + method); - }); - } - - private static final class FakeDeviceBuilder implements CuVSMatrix.Builder { - private final CuVSDeviceMatrix dataset; - private final RuntimeException closeFailure; - private final AtomicInteger closeCalls = new AtomicInteger(); - - private FakeDeviceBuilder(CuVSDeviceMatrix dataset, RuntimeException closeFailure) { - this.dataset = dataset; - this.closeFailure = closeFailure; - } - - @Override - public void addVector(float[] vector) {} - - @Override - public void addVector(byte[] vector) {} - - @Override - public void addVector(int[] vector) {} - - @Override - public void addVector(short[] vector) {} - - @Override - public CuVSDeviceMatrix build() { - return dataset; - } - - @Override - public void close() { - closeCalls.incrementAndGet(); - if (closeFailure != null) { - throw closeFailure; - } - } - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java index 4ac27b988b..3b05beaacf 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java @@ -12,11 +12,7 @@ import com.carrotsearch.randomizedtesting.annotations.Name; import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; import java.util.Arrays; -import java.util.HashSet; -import java.util.LinkedHashMap; import java.util.List; -import java.util.Map; -import java.util.Set; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; @@ -30,9 +26,6 @@ import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.index.NoMergePolicy; -import org.apache.lucene.index.Term; -import org.apache.lucene.index.TieredMergePolicy; import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.store.ByteBuffersDirectory; import org.apache.lucene.store.Directory; @@ -82,8 +75,7 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { } try (Directory dir = newDirectory(new ByteBuffersDirectory()); - IndexWriter w = - new IndexWriter(dir, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))) { + IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { for (int i = 0; i < R; i++) { Document doc = new Document(); doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); @@ -100,7 +92,6 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { assertEquals(1, subReaders.get(i).reader().getFloatVectorValues(F).size()); } } - w.getConfig().setMergePolicy(new TieredMergePolicy()); w.forceMerge(1); try (DirectoryReader reader = DirectoryReader.open(w)) { @@ -175,148 +166,6 @@ public void testCosineSimilarity() throws Exception { } } - public void testForceMergeUsesOnlyLiveSparseVectors() throws Exception { - final String vectorField = "vector"; - final int dimensions = 129; - Map expected = new LinkedHashMap<>(); - - try (Directory directory = newDirectory(new ByteBuffersDirectory())) { - try (IndexWriter writer = - new IndexWriter( - directory, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))) { - for (int segment = 0; segment < 3; segment++) { - for (int row = 0; row < 5; row++) { - String id = segment + "-" + row; - Document document = new Document(); - document.add(new StringField("id", id, Field.Store.YES)); - if (row < 4) { - float[] vector = deterministicVector(segment * 5 + row, dimensions); - document.add(new KnnFloatVectorField(vectorField, vector, EUCLIDEAN)); - if (row != 1) { - expected.put(id, vector); - } - } - writer.addDocument(document); - } - writer.commit(); - } - for (int segment = 0; segment < 3; segment++) { - writer.deleteDocuments(new Term("id", segment + "-1")); - } - writer.commit(); - - try (DirectoryReader sourceReader = DirectoryReader.open(writer)) { - assertEquals("the test requires three source segments", 3, sourceReader.leaves().size()); - for (LeafReaderContext context : sourceReader.leaves()) { - LeafReader sourceLeaf = context.reader(); - assertTrue("each source segment must carry a deletion", sourceLeaf.hasDeletions()); - assertEquals(5, sourceLeaf.maxDoc()); - assertEquals(4, sourceLeaf.numDocs()); - assertEquals(4, sourceLeaf.getFloatVectorValues(vectorField).size()); - } - } - - writer.getConfig().setMergePolicy(new TieredMergePolicy()); - writer.forceMerge(1); - } - - TestUtil.checkIndex(directory); - try (DirectoryReader reader = DirectoryReader.open(directory)) { - LeafReader leaf = getOnlyLeafReader(reader); - FloatVectorValues values = leaf.getFloatVectorValues(vectorField); - assertNotNull(values); - assertEquals(expected.size(), values.size()); - - Set seen = new HashSet<>(); - for (int ordinal = 0; ordinal < values.size(); ordinal++) { - int documentId = values.ordToDoc(ordinal); - String id = leaf.storedFields().document(documentId).get("id"); - assertTrue("Unexpected or duplicate vector for " + id, seen.add(id)); - assertArrayEquals(expected.get(id), values.vectorValue(ordinal), 0.0f); - } - assertEquals(expected.keySet(), seen); - - for (Map.Entry entry : expected.entrySet()) { - var hits = - leaf.searchNearestVectors( - vectorField, entry.getValue(), expected.size(), null, 1_000); - boolean found = false; - for (var hit : hits.scoreDocs) { - found |= entry.getKey().equals(leaf.storedFields().document(hit.doc).get("id")); - } - assertTrue("Exact vector was not searchable for " + entry.getKey(), found); - } - } - } - } - - public void testForceMergeWithZeroLiveVectors() throws Exception { - assertTrivialLiveVectorMerge(0); - } - - public void testForceMergeWithOneLiveVector() throws Exception { - assertTrivialLiveVectorMerge(1); - } - - private void assertTrivialLiveVectorMerge(int liveVectors) throws Exception { - final String vectorField = "vector"; - final int dimensions = 129; - try (Directory directory = newDirectory(new ByteBuffersDirectory())) { - try (IndexWriter writer = - new IndexWriter( - directory, newIndexWriterConfig().setMergePolicy(NoMergePolicy.INSTANCE))) { - for (int id = 0; id < 3; id++) { - Document vectorDocument = new Document(); - vectorDocument.add(new StringField("id", "vector-" + id, Field.Store.YES)); - vectorDocument.add( - new KnnFloatVectorField(vectorField, deterministicVector(id, dimensions), EUCLIDEAN)); - writer.addDocument(vectorDocument); - - Document sparseDocument = new Document(); - sparseDocument.add(new StringField("id", "sparse-" + id, Field.Store.YES)); - writer.addDocument(sparseDocument); - writer.commit(); - } - for (int id = liveVectors; id < 3; id++) { - writer.deleteDocuments(new Term("id", "vector-" + id)); - } - writer.commit(); - - try (DirectoryReader sourceReader = DirectoryReader.open(writer)) { - assertEquals("the test requires three source segments", 3, sourceReader.leaves().size()); - for (LeafReaderContext context : sourceReader.leaves()) { - assertEquals(1, context.reader().getFloatVectorValues(vectorField).size()); - } - } - - writer.getConfig().setMergePolicy(new TieredMergePolicy()); - writer.forceMerge(1); - } - - TestUtil.checkIndex(directory); - try (DirectoryReader reader = DirectoryReader.open(directory)) { - LeafReader leaf = getOnlyLeafReader(reader); - FloatVectorValues values = leaf.getFloatVectorValues(vectorField); - if (liveVectors == 0) { - assertTrue(values == null || values.size() == 0); - } else { - assertNotNull(values); - assertEquals(1, values.size()); - assertEquals("vector-0", leaf.storedFields().document(values.ordToDoc(0)).get("id")); - assertArrayEquals(deterministicVector(0, dimensions), values.vectorValue(0), 0.0f); - } - } - } - } - - private static float[] deterministicVector(int id, int dimensions) { - float[] vector = new float[dimensions]; - for (int dimension = 0; dimension < dimensions; dimension++) { - vector[dimension] = id * 10.0f + dimension * 0.01f; - } - return vector; - } - @Override protected VectorEncoding randomVectorEncoding() { return VectorEncoding.FLOAT32; diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java index 89011fc2ef..6aead1bd8d 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java @@ -4,10 +4,7 @@ */ package com.nvidia.cuvs.lucene; -import com.nvidia.cuvs.CuVSMatrix; import java.io.IOException; -import java.lang.reflect.Proxy; -import java.util.concurrent.atomic.AtomicInteger; import org.apache.lucene.tests.util.LuceneTestCase; import org.junit.Test; @@ -50,182 +47,4 @@ public void testHandleThrowableWrapsCheckedExceptionWithCause() { assertSame(exception, thrown.getCause()); } - - @Test - public void testMatrixBuilderCloseFailureClosesTransferredMatrix() { - IllegalStateException builderFailure = new IllegalStateException("builder close"); - AtomicInteger matrixCloses = new AtomicInteger(); - CuVSMatrix matrix = fakeMatrix(matrixCloses, null); - FakeMatrixBuilder builder = new FakeMatrixBuilder(matrix, builderFailure); - - IllegalStateException thrown = - assertThrows( - IllegalStateException.class, - () -> MatrixBuilderLifecycle.build(builder, CuVSMatrix.Builder::build)); - - assertSame(builderFailure, thrown); - assertEquals(1, builder.closeCalls.get()); - assertEquals(1, matrixCloses.get()); - } - - @Test - public void testMatrixCloseFailureIsSuppressedOnBuilderCloseFailure() { - IllegalStateException builderFailure = new IllegalStateException("builder close"); - IllegalArgumentException matrixFailure = new IllegalArgumentException("matrix close"); - AtomicInteger matrixCloses = new AtomicInteger(); - CuVSMatrix matrix = fakeMatrix(matrixCloses, matrixFailure); - FakeMatrixBuilder builder = new FakeMatrixBuilder(matrix, builderFailure); - - IllegalStateException thrown = - assertThrows( - IllegalStateException.class, - () -> MatrixBuilderLifecycle.build(builder, CuVSMatrix.Builder::build)); - - assertSame(builderFailure, thrown); - assertArrayEquals(new Throwable[] {matrixFailure}, thrown.getSuppressed()); - assertEquals(1, matrixCloses.get()); - } - - @Test - public void testMatrixOperationFailureRemainsPrimaryWhenBuilderCloseFails() { - IOException operationFailure = new IOException("populate"); - IllegalStateException builderFailure = new IllegalStateException("builder close"); - FakeMatrixBuilder builder = new FakeMatrixBuilder(null, builderFailure); - - IOException thrown = - assertThrows( - IOException.class, - () -> - MatrixBuilderLifecycle.build( - builder, - ignored -> { - throw operationFailure; - })); - - assertSame(operationFailure, thrown); - assertArrayEquals(new Throwable[] {builderFailure}, thrown.getSuppressed()); - assertEquals(1, builder.closeCalls.get()); - } - - @Test - public void testOwnedIndexClosesUntransferredDatasetOnce() throws Exception { - TrackingCloseable dataset = new TrackingCloseable(null); - Utils.OwnedIndex owned = Utils.ownDataset(dataset); - - owned.close(); - owned.close(); - - assertEquals(1, dataset.closeCount); - } - - @Test - public void testOwnedIndexDoesNotDirectlyCloseDatasetAfterSuccessfulIndexClose() - throws Exception { - TrackingCloseable dataset = new TrackingCloseable(null); - TrackingCloseable index = new TrackingCloseable(null); - Utils.OwnedIndex owned = Utils.ownDataset(dataset); - owned.transferTo(index); - - owned.close(); - - assertEquals(1, index.closeCount); - assertEquals(0, dataset.closeCount); - } - - @Test - public void testOwnedIndexPreservesBodyAndNestedCleanupFailures() { - IOException bodyFailure = new IOException("body"); - IOException indexCloseFailure = new IOException("index close"); - IOException datasetCloseFailure = new IOException("dataset close"); - TrackingCloseable dataset = new TrackingCloseable(datasetCloseFailure); - TrackingCloseable index = new TrackingCloseable(indexCloseFailure); - - IOException thrown = - assertThrows( - IOException.class, - () -> { - try (Utils.OwnedIndex owned = Utils.ownDataset(dataset)) { - owned.transferTo(index); - throw bodyFailure; - } - }); - - assertSame(bodyFailure, thrown); - assertEquals(1, thrown.getSuppressed().length); - assertSame(indexCloseFailure, thrown.getSuppressed()[0]); - assertEquals(1, indexCloseFailure.getSuppressed().length); - assertSame(datasetCloseFailure, indexCloseFailure.getSuppressed()[0]); - assertEquals(1, index.closeCount); - assertEquals(1, dataset.closeCount); - } - - private static final class TrackingCloseable implements AutoCloseable { - private final Exception failure; - private int closeCount; - - private TrackingCloseable(Exception failure) { - this.failure = failure; - } - - @Override - public void close() throws Exception { - closeCount++; - if (failure != null) { - throw failure; - } - } - } - - private static CuVSMatrix fakeMatrix(AtomicInteger closeCalls, RuntimeException closeFailure) { - return (CuVSMatrix) - Proxy.newProxyInstance( - CuVSMatrix.class.getClassLoader(), - new Class[] {CuVSMatrix.class}, - (proxy, method, args) -> { - if (method.getName().equals("close") && method.getParameterCount() == 0) { - closeCalls.incrementAndGet(); - if (closeFailure != null) { - throw closeFailure; - } - return null; - } - throw new AssertionError("Unexpected matrix call: " + method); - }); - } - - private static final class FakeMatrixBuilder implements CuVSMatrix.Builder { - private final CuVSMatrix matrix; - private final RuntimeException closeFailure; - private final AtomicInteger closeCalls = new AtomicInteger(); - - private FakeMatrixBuilder(CuVSMatrix matrix, RuntimeException closeFailure) { - this.matrix = matrix; - this.closeFailure = closeFailure; - } - - @Override - public void addVector(float[] vector) {} - - @Override - public void addVector(byte[] vector) {} - - @Override - public void addVector(int[] vector) {} - - @Override - public void addVector(short[] vector) {} - - @Override - public CuVSMatrix build() { - return matrix; - } - - @Override - public void close() { - closeCalls.incrementAndGet(); - if (closeFailure != null) { - throw closeFailure; - } - } - } } From f543ff06c6223a0cd4e712c1d11a0228c91e25b9 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 18 Sep 2026 19:15:03 +0000 Subject: [PATCH 07/21] Validate accelerated HNSW merge replay counts --- .../Lucene99AcceleratedHNSWVectorsWriter.java | 29 ++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 863877a9e3..9bc4659b2e 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -306,10 +306,32 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws CuVSMatrix.Builder builder = CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT); KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); + int replayed = 0; + boolean hasExtraVector = false; for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { + if (replayed == size) { + hasExtraVector = true; + break; + } builder.addVector(mergedVectors.vectorValue(it.index())); + replayed = Math.incrementExact(replayed); } CuVSHostMatrix dataset = builder.build(); + if (hasExtraVector || replayed != size) { + int observed = hasExtraVector ? Math.incrementExact(replayed) : replayed; + IOException mismatch = + new IOException( + "Merged vector count changed between passes: expected " + + size + + (hasExtraVector ? ", observed at least " : ", observed ") + + observed); + try { + dataset.close(); + } catch (Throwable closeFailure) { + mismatch.addSuppressed(closeFailure); + } + throw mismatch; + } writeNonTrivialField(fieldInfo, dataset); } catch (Throwable t) { Utils.handleThrowable(t); @@ -324,7 +346,12 @@ private static int countMergedVectors(FieldInfo fieldInfo, MergeState mergeState int count = 0; KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { - count++; + try { + count = Math.incrementExact(count); + } catch (ArithmeticException tooManyVectors) { + throw new IOException( + "Merged vector count exceeds the supported integer range", tooManyVectors); + } } return count; } From f3bef0ece073060fbc83dfabba6b806629c0526d Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Fri, 18 Sep 2026 20:07:59 +0000 Subject: [PATCH 08/21] Close matrix builders during merge replay --- .../java-api-com-nvidia-cuvs-cuvsmatrix.md | 59 +++++-- ...nvidia-cuvs-lucene-acceleratedhnswutils.md | 29 ++-- ...ne-lucene99acceleratedhnswvectorswriter.md | 16 +- ...leratedhnswbinaryquantizedvectorswriter.md | 10 +- ...leratedhnswscalarquantizedvectorswriter.md | 10 +- ...lucene-api-com-nvidia-cuvs-lucene-utils.md | 81 ++++------ fern/pages/other/multidimensional_arrays.md | 12 +- .../main/java/com/nvidia/cuvs/CuVSMatrix.java | 20 ++- .../com/nvidia/cuvs/spi/JDKProvider.java | 38 ++++- .../java/com/nvidia/cuvs/CuVSMatrixIT.java | 56 +++++++ .../Lucene99AcceleratedHNSWVectorsWriter.java | 51 +++--- .../java/com/nvidia/cuvs/lucene/Utils.java | 117 +++----------- .../TestAcceleratedHNSWMergeReplay.java | 147 ++++++++++++++++++ 13 files changed, 419 insertions(+), 227 deletions(-) create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMergeReplay.java diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsmatrix.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsmatrix.md index 29e3e618fd..f07a620ffd 100644 --- a/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsmatrix.md +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsmatrix.md @@ -114,6 +114,35 @@ Adds a single vector to the matrix. Each element is a raw float16 bit pattern st _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:102`_ +### build + +```java +T build() +``` + +Completes the matrix and transfers ownership to the caller. + +If this method fails, callers should close the builder. Built-in builders then release +matrix storage allocated during builder construction; providers that inherit the default +no-op `#close()` implementation do not. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:111`_ + +### close + +```java +@Override default void close() +``` + +Closes this builder. Built-in builders release matrix storage unless ownership was +transferred by a successful `#build()`. + +The default implementation preserves compatibility with providers compiled before +builders became closeable. Builders that allocate storage before `#build()` should +override this method. + +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:121`_ + ### hostBuilder ```java @@ -134,7 +163,7 @@ Returns a builder to create a new instance of a host-memory matrix a builder for creating a `CuVSHostMatrix` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:115`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:133`_ ### hostBuilder @@ -158,7 +187,7 @@ Returns a builder to create a new instance of a host-memory matrix a builder for creating a `CuVSDeviceMatrix` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:129`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:147`_ ### deviceBuilder @@ -181,7 +210,7 @@ Returns a builder to create a new instance of a dataset a builder for creating a `CuVSDeviceMatrix` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:144`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:162`_ ### deviceBuilder @@ -206,7 +235,7 @@ Returns a builder to create a new instance of a dataset a builder for creating a `CuVSDeviceMatrix` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:160`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:178`_ ### size @@ -220,7 +249,7 @@ Gets the size of the dataset Size of the dataset -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:176`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:194`_ ### columns @@ -235,7 +264,7 @@ or the graph degree for the graph represented as a list of neighbours Dimensions of the vectors in the dataset -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:184`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:202`_ ### dataType @@ -249,7 +278,7 @@ Gets the element type a `DataType` describing the matrix element type -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:191`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:209`_ ### getRow @@ -265,7 +294,7 @@ Get a view (0-copy) of the row data, as a list of integers (32 bit) | --- | --- | | `row` | the row for which to return the data | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:198`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:216`_ ### toArray @@ -281,7 +310,7 @@ Copies the content of this dataset to an on-heap Java matrix (array of arrays). | --- | --- | | `array` | the destination array. Must be of length `CuVSMatrix#size()` or bigger, and each element must be of length `CuVSMatrix#columns()` or bigger. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:206`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:224`_ ### toArray @@ -297,7 +326,7 @@ Copies the content of this dataset to an on-heap Java matrix (array of arrays). | --- | --- | | `array` | the destination array. Must be of length `CuVSMatrix#size()` or bigger, and each element must be of length `CuVSMatrix#columns()` or bigger. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:214`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:232`_ ### toArray @@ -313,7 +342,7 @@ Copies the content of this dataset to an on-heap Java matrix (array of arrays). | --- | --- | | `array` | the destination array. Must be of length `CuVSMatrix#size()` or bigger, and each element must be of length `CuVSMatrix#columns()` or bigger. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:222`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:240`_ ### toHost @@ -331,7 +360,7 @@ same element type and dimension. | --- | --- | | `hostMatrix` | the host-memory-backed matrix to fill. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:231`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:249`_ ### toHost @@ -345,7 +374,7 @@ the device matrix. The returned host matrix will need to be managed by the caller, which will be responsible to call `CuVSMatrix#close()` to free its resources when done. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:240`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:258`_ ### toDevice @@ -363,7 +392,7 @@ same element type and dimension. | --- | --- | | `deviceMatrix` | the device-memory-backed matrix to fill. | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:249`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:267`_ ### toDevice @@ -377,6 +406,6 @@ the host matrix. The returned device matrix will need to be managed by the caller, which will be responsible to call `CuVSMatrix#close()` to free its resources when done. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:258`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:276`_ _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java:17`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md index 219149bc6d..b3bd116a4c 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md @@ -21,7 +21,7 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens Creates a dummy HNSW graph for a single vector. The graph will have 1 level with 1 node and no neighbors. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:54`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:55`_ ### createMultiLayerHnswGraph @@ -35,7 +35,18 @@ M = ceil(cagraGraphDegree / 2), where cagraGraphDegree is the CAGRA adjacency li Each layer contains 1/M nodes from the previous layer Creates layers until the highest layer has ≤ M nodes -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:81`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:82`_ + +### createMultiLayerHnswGraph + +```java +static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, int dimensions, CuVSMatrix adjacencyListMatrix, CuVSMatrix vectorDataset, int hnswLayers, CagraIndexParams params, QuantizationType quantization) throws Throwable +``` + +Creates a multi-layer HNSW graph from a native matrix without copying the complete dataset to +the Java heap. The list view copies only rows selected for an upper layer. + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:176`_ ### writeGraph @@ -62,7 +73,7 @@ a 2D array of offsets | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:237`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:288`_ ### writeMeta @@ -91,7 +102,7 @@ Writes the meta information for the index. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:302`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:353`_ ### printInfoStream @@ -107,7 +118,7 @@ A utility method to print info/debugging messages using InfoStream. | --- | --- | | `msg` | the debugging message to print | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:384`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:435`_ ### writeEmpty @@ -129,7 +140,7 @@ Writes an empty meta information for the field. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:396`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:447`_ ### quantizeFloatVectorsToBinary @@ -152,7 +163,7 @@ Bits are packed: 8 dimensions per byte. A list of byte binary representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:409`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:460`_ ### quantizeFloatVectorsToScalar @@ -172,6 +183,6 @@ Scalar quantization. A list of byte scalar representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:451`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:502`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:31`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:32`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md index 2f30caeb07..7f5698a6a3 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md @@ -37,7 +37,7 @@ Initializes `Lucene99AcceleratedHNSWVectorsWriter` | --- | --- | | `IOException` | IOException | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:86`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:89`_ ### addField @@ -47,7 +47,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:128`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:131`_ ### flush @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:203`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:222`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:291`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:364`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:300`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:373`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:320`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:393`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:330`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:403`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:52`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:55`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md index 91a4e67303..83a7741bed 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:215`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:213`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:302`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:300`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:333`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:331`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:353`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:351`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:362`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:360`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:57`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md index 9c51cb26e9..66d4678107 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:241`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:239`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:326`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:324`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:357`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:355`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:377`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:375`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:386`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:384`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:56`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md index 0d2a9cb1a6..5717013e42 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md @@ -45,7 +45,7 @@ never returns; always throws | --- | --- | | `IOException` | | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:43`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:39`_ ### createFloatMatrix @@ -70,54 +70,51 @@ without creating intermediate heap arrays. an instance of CuVSMatrix -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:63`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:59`_ -### createByteMatrix +### createHostFloatMatrix ```java -static CuVSMatrix createByteMatrix( List data, int bytesPerVector, CuVSResources resources) +static CuVSHostMatrix createHostFloatMatrix(List data, int dimensions) ``` -A method to build a CuVSMatrix from a list of byte vectors (for binary quantized vectors). +Builds a host-memory CuVSMatrix from a list of float vectors. -Uses CuVSMatrix.Builder to copy vectors directly to device memory -without creating intermediate heap arrays. +Copies vectors directly into native host memory without creating an intermediate \{@code +float[][]\} on the heap. **Parameters** | Name | Description | | --- | --- | -| `data` | The byte vectors (packed bits for binary quantization) | -| `bytesPerVector` | The number of bytes in each vector | -| `resources` | The CuVS resources for device matrix creation | +| `data` | The float vectors | +| `dimensions` | The number of float elements in each vector | **Returns** -an instance of CuVSMatrix with BYTE data type +a host-memory CuVSMatrix -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:92`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:87`_ -### createByteMatrixFromArray +### createHostByteMatrix ```java -static CuVSMatrix createByteMatrixFromArray( byte[][] data, int bytesPerVector, CuVSResources resources) +static CuVSHostMatrix createHostByteMatrix(List data, int bytesPerVector) ``` -A method to build a CuVSMatrix from a 2D byte array (for binary quantized vectors). +Builds a host-memory CuVSMatrix from a list of byte vectors. -**Parameters** +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:98`_ -| Name | Description | -| --- | --- | -| `data` | The 2D byte array (packed bits for binary quantization) | -| `bytesPerVector` | The number of bytes in each vector | -| `resources` | The CuVS resources for device matrix creation | +### createHostByteMatrixFromArray -**Returns** +```java +static CuVSHostMatrix createHostByteMatrixFromArray(byte[][] data, int bytesPerVector) +``` -an instance of CuVSMatrix with BYTE data type +Builds a host-memory CuVSMatrix from a 2D byte array. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:119`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:109`_ ### nanosToMillis @@ -137,7 +134,7 @@ A utility method to convert nanoseconds to milliseconds. milliseconds -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:141`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:125`_ ### cuVSResourcesOrNull @@ -151,7 +148,7 @@ Creates an instance of CuVSResources. an instance of CuVSResources -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:150`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:134`_ ### handleThrowableWithIgnore @@ -174,33 +171,7 @@ A utility method that conditionally ignores certain throwable objects | --- | --- | | `IOException` | | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:178`_ - -### createListFromMergedVectors - -```java -static List createListFromMergedVectors(FloatVectorValues mergedVectorValues) throws IOException -``` - -Creates a list of float vectors from the input - -**Parameters** - -| Name | Description | -| --- | --- | -| `mergedVectorValues` | instance of `FloatVectorValues` | - -**Returns** - -a list of float arrays - -**Throws** - -| Type | Description | -| --- | --- | -| `IOException` | I/O Exception | - -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:192`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:162`_ ### info @@ -218,6 +189,6 @@ Utility to print info/debug messages via InfoStream. | `component` | the name of the index writer | | `msg` | the log message to push via the InfoStream | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:210`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:176`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:26`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:22`_ diff --git a/fern/pages/other/multidimensional_arrays.md b/fern/pages/other/multidimensional_arrays.md index c5f7ec7d5e..072afa1087 100644 --- a/fern/pages/other/multidimensional_arrays.md +++ b/fern/pages/other/multidimensional_arrays.md @@ -165,16 +165,16 @@ import com.nvidia.cuvs.CuVSDeviceMatrix; long nRows = 100_000; long nFeatures = 128; -try (CuVSResources resources = CuVSResources.create()) { - CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, nRows, nFeatures, CuVSMatrix.DataType.FLOAT); - +try (CuVSResources resources = CuVSResources.create(); + CuVSMatrix.Builder builder = + CuVSMatrix.deviceBuilder( + resources, nRows, nFeatures, CuVSMatrix.DataType.FLOAT)) { for (long row = 0; row < nRows; row++) { builder.addVector(loadVector(row)); } - try (CuVSMatrix dataset = builder.build()) { + // A successful build transfers matrix ownership from the builder. + try (CuVSDeviceMatrix dataset = builder.build()) { // Pass dataset to NVIDIA cuVS Java APIs. } } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java index 097efb1003..bc45679670 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSMatrix.java @@ -73,7 +73,7 @@ static CuVSMatrix ofArray(byte[][] vectors) { * A builder to construct a new matrix one row at a time * @param the CuVSMatrix type to build */ - interface Builder { + interface Builder extends AutoCloseable { /** * Adds a single vector to the matrix. * @@ -102,7 +102,25 @@ interface Builder { */ void addVector(short[] vector); + /** + * Completes the matrix and transfers ownership to the caller. + * + *

If this method fails, callers should close the builder. Built-in builders then release + * matrix storage allocated during builder construction; providers that inherit the default + * no-op {@link #close()} implementation do not. + */ T build(); + + /** + * Closes this builder. Built-in builders release matrix storage unless ownership was + * transferred by a successful {@link #build()}. + * + *

The default implementation preserves compatibility with providers compiled before + * builders became closeable. Builders that allocate storage before {@link #build()} should + * override this method. + */ + @Override + default void close() {} } /** diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 1f16a4e904..c240e35d3e 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -617,7 +617,8 @@ public CuVSMatrix newMatrixFromArray(byte[][] vectors) { return dataset; } - private abstract static class MatrixBuilder { + private abstract static class MatrixBuilder + implements AutoCloseable { protected final long columns; protected final long size; @@ -626,6 +627,7 @@ private abstract static class MatrixBuilder { protected final long rowSize; protected final long rowBytes; protected int currentRow; + private boolean closed; protected MatrixBuilder(T matrix, long size, long columns) { this.columns = columns; @@ -635,6 +637,7 @@ protected MatrixBuilder(T matrix, long size, long columns) { this.rowSize = columns * elementSize; this.rowBytes = rowSize; this.currentRow = 0; + this.closed = false; } protected MatrixBuilder(T matrix, long size, long columns, int rowStride) { @@ -646,9 +649,11 @@ protected MatrixBuilder(T matrix, long size, long columns, int rowStride) { this.rowBytes = columns * elementSize; this.currentRow = 0; + this.closed = false; } public void addVector(float[] vector) { + ensureOpen(); if (vector.length != columns) { throw new IllegalArgumentException( String.format( @@ -658,6 +663,7 @@ public void addVector(float[] vector) { } public void addVector(byte[] vector) { + ensureOpen(); if (vector.length != columns) { throw new IllegalArgumentException( String.format( @@ -667,6 +673,7 @@ public void addVector(byte[] vector) { } public void addVector(int[] vector) { + ensureOpen(); if (vector.length != columns) { throw new IllegalArgumentException( String.format( @@ -676,6 +683,7 @@ public void addVector(int[] vector) { } public void addVector(short[] vector) { + ensureOpen(); if (vector.length != columns) { throw new IllegalArgumentException( String.format( @@ -684,6 +692,27 @@ public void addVector(short[] vector) { internalAddVector(MemorySegment.ofArray(vector)); } + protected final T transferOwnership() { + ensureOpen(); + closed = true; + return matrix; + } + + protected final void ensureOpen() { + if (closed) { + throw new IllegalStateException("matrix builder is closed"); + } + } + + @Override + public final void close() { + if (closed) { + return; + } + closed = true; + matrix.close(); + } + protected abstract void internalAddVector(MemorySegment vector); } @@ -772,11 +801,12 @@ private void flushBuffer(MemorySegment hostBuffer) { @Override public CuVSDeviceMatrix build() { + ensureOpen(); try (var access = resources.access()) { var hostBuffer = CuVSResourcesImpl.getHostBuffer(access); flushBuffer(hostBuffer); } - return matrix; + return transferOwnership(); } } @@ -831,7 +861,7 @@ protected void internalAddVector(MemorySegment vector) { @Override public CuVSDeviceMatrix build() { - return matrix; + return transferOwnership(); } } @@ -861,7 +891,7 @@ protected void internalAddVector(MemorySegment vector) { @Override public CuVSHostMatrix build() { - return matrix; + return transferOwnership(); } } } diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java index 954a7d1e4f..4f369b207c 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSMatrixIT.java @@ -7,6 +7,7 @@ import static com.carrotsearch.randomizedtesting.RandomizedTest.*; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; import com.carrotsearch.randomizedtesting.RandomizedRunner; import com.nvidia.cuvs.spi.CuVSProvider; @@ -309,6 +310,61 @@ public void testFloatDatasetDeviceBuilder() throws Throwable { } } + @Test + public void testClosingHostBuilderReleasesUnbuiltMatrix() { + assertClosingBuilderReleasesUnbuiltMatrix( + CuVSMatrix.hostBuilder(2, 4, CuVSMatrix.DataType.FLOAT)); + } + + @Test + public void testClosingDeviceBuilderReleasesUnbuiltMatrix() throws Throwable { + try (var resources = CheckedCuVSResources.create()) { + assertClosingBuilderReleasesUnbuiltMatrix( + CuVSMatrix.deviceBuilder(resources, 2, 4, CuVSMatrix.DataType.FLOAT)); + } + } + + private static void assertClosingBuilderReleasesUnbuiltMatrix(CuVSMatrix.Builder builder) { + builder.addVector(new float[4]); + builder.close(); + builder.close(); + + assertThrows(IllegalStateException.class, builder::build); + assertThrows(IllegalStateException.class, () -> builder.addVector(new float[4])); + } + + @Test + public void testClosingHostBuilderAfterBuildDoesNotCloseMatrix() { + assertClosingBuilderAfterBuildDoesNotCloseMatrix( + CuVSMatrix.hostBuilder(1, 4, CuVSMatrix.DataType.FLOAT)); + } + + @Test + public void testClosingDeviceBuilderAfterBuildDoesNotCloseMatrix() throws Throwable { + try (var resources = CheckedCuVSResources.create()) { + assertClosingBuilderAfterBuildDoesNotCloseMatrix( + CuVSMatrix.deviceBuilder(resources, 1, 4, CuVSMatrix.DataType.FLOAT)); + } + } + + private static void assertClosingBuilderAfterBuildDoesNotCloseMatrix( + CuVSMatrix.Builder builder) { + CuVSMatrix matrix; + try (builder) { + builder.addVector(new float[] {1f, 2f, 3f, 4f}); + matrix = builder.build(); + + assertThrows(IllegalStateException.class, builder::build); + assertThrows(IllegalStateException.class, () -> builder.addVector(new float[4])); + } + + try (matrix) { + float[][] actual = new float[1][4]; + matrix.toArray(actual); + assertArrayEquals(new float[] {1f, 2f, 3f, 4f}, actual[0], DELTA); + } + } + private void testIntDatasetBuilder(int rows, int cols, CuVSMatrix.Builder builder) { var data = new int[rows][cols]; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 9bc4659b2e..bbfd4309b3 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -303,41 +303,44 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws FloatVectorValues mergedVectors = KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); int dims = fieldInfo.getVectorDimension(); - CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT); + CuVSHostMatrix dataset = + buildMergedDataset( + mergedVectors, size, CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT)); + writeNonTrivialField(fieldInfo, dataset); + } catch (Throwable t) { + Utils.handleThrowable(t); + } + } + + /* Replays merged vectors into a builder that remains responsible for storage until build. */ + static CuVSHostMatrix buildMergedDataset( + FloatVectorValues mergedVectors, int expectedSize, CuVSMatrix.Builder builder) + throws IOException { + try (builder) { KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); int replayed = 0; - boolean hasExtraVector = false; for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { - if (replayed == size) { - hasExtraVector = true; - break; + if (replayed == expectedSize) { + throw mergeReplayMismatch(expectedSize, (long) replayed + 1L, true); } builder.addVector(mergedVectors.vectorValue(it.index())); replayed = Math.incrementExact(replayed); } - CuVSHostMatrix dataset = builder.build(); - if (hasExtraVector || replayed != size) { - int observed = hasExtraVector ? Math.incrementExact(replayed) : replayed; - IOException mismatch = - new IOException( - "Merged vector count changed between passes: expected " - + size - + (hasExtraVector ? ", observed at least " : ", observed ") - + observed); - try { - dataset.close(); - } catch (Throwable closeFailure) { - mismatch.addSuppressed(closeFailure); - } - throw mismatch; + if (replayed != expectedSize) { + throw mergeReplayMismatch(expectedSize, replayed, false); } - writeNonTrivialField(fieldInfo, dataset); - } catch (Throwable t) { - Utils.handleThrowable(t); + return builder.build(); } } + private static IOException mergeReplayMismatch(int expected, long observed, boolean lowerBound) { + return new IOException( + "Merged vector count changed between passes: expected " + + expected + + (lowerBound ? ", observed at least " : ", observed ") + + observed); + } + /** Counts the live vectors that the merge iterator will actually yield. */ private static int countMergedVectors(FieldInfo fieldInfo, MergeState mergeState) throws IOException { diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java index 91bbc38e2c..97bc646225 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -4,18 +4,14 @@ */ package com.nvidia.cuvs.lucene; -import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; - +import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.CuVSResources; import java.io.IOException; import java.time.Duration; -import java.util.ArrayList; import java.util.List; import java.util.logging.Level; import java.util.logging.Logger; -import org.apache.lucene.index.FloatVectorValues; -import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.util.InfoStream; /** @@ -89,87 +85,36 @@ static CuVSMatrix createFloatMatrix(List data, int dimensions, CuVSReso * @param dimensions The number of float elements in each vector * @return a host-memory CuVSMatrix */ - static CuVSMatrix createHostFloatMatrix(List data, int dimensions) { - CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT); - for (float[] vector : data) { - builder.addVector(vector); - } - return builder.build(); - } - - /** - * A method to build a CuVSMatrix from a list of byte vectors (for binary quantized vectors). - * - * Uses CuVSMatrix.Builder to copy vectors directly to device memory - * without creating intermediate heap arrays. - * - * @param data The byte vectors (packed bits for binary quantization) - * @param bytesPerVector The number of bytes in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix with BYTE data type - */ - static CuVSMatrix createByteMatrix( - List data, int bytesPerVector, CuVSResources resources) { - // Use Builder pattern to avoid intermediate byte[][] allocation - // and copy directly from List to device memory - CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.size(), // rows (number of vectors) - bytesPerVector, // columns (bytes per vector) - CuVSMatrix.DataType.BYTE); - - // Add vectors one by one - builder copies directly to device memory - for (byte[] vector : data) { - builder.addVector(vector); + static CuVSHostMatrix createHostFloatMatrix(List data, int dimensions) { + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT)) { + for (float[] vector : data) { + builder.addVector(vector); + } + return builder.build(); } - - return builder.build(); } /** Builds a host-memory CuVSMatrix from a list of byte vectors. */ - static CuVSMatrix createHostByteMatrix(List data, int bytesPerVector) { - CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE); - for (byte[] vector : data) { - builder.addVector(vector); - } - return builder.build(); - } - - /** - * A method to build a CuVSMatrix from a 2D byte array (for binary quantized vectors). - * - * @param data The 2D byte array (packed bits for binary quantization) - * @param bytesPerVector The number of bytes in each vector - * @param resources The CuVS resources for device matrix creation - * @return an instance of CuVSMatrix with BYTE data type - */ - static CuVSMatrix createByteMatrixFromArray( - byte[][] data, int bytesPerVector, CuVSResources resources) { - CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.length, // rows (number of vectors) - bytesPerVector, // columns (bytes per vector) - CuVSMatrix.DataType.BYTE); - - // Add vectors one by one - builder copies directly to device memory - for (byte[] vector : data) { - builder.addVector(vector); + static CuVSHostMatrix createHostByteMatrix(List data, int bytesPerVector) { + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE)) { + for (byte[] vector : data) { + builder.addVector(vector); + } + return builder.build(); } - return builder.build(); } /** Builds a host-memory CuVSMatrix from a 2D byte array. */ - static CuVSMatrix createHostByteMatrixFromArray(byte[][] data, int bytesPerVector) { - CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE); - for (byte[] vector : data) { - builder.addVector(vector); + static CuVSHostMatrix createHostByteMatrixFromArray(byte[][] data, int bytesPerVector) { + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(data.length, bytesPerVector, CuVSMatrix.DataType.BYTE)) { + for (byte[] vector : data) { + builder.addVector(vector); + } + return builder.build(); } - return builder.build(); } /** @@ -222,24 +167,6 @@ static void handleThrowableWithIgnore(Throwable t, String msg) throws IOExceptio handleThrowable(t); } - /** - * Creates a list of float vectors from the input - * - * @param mergedVectorValues instance of {@link FloatVectorValues} - * @return a list of float arrays - * @throws IOException I/O Exception - */ - static List createListFromMergedVectors(FloatVectorValues mergedVectorValues) - throws IOException { - List vectors = new ArrayList(); - KnnVectorValues.DocIndexIterator iter = mergedVectorValues.iterator(); - for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) { - float[] vector = mergedVectorValues.vectorValue(iter.index()); - vectors.add(vector.clone()); - } - return vectors; - } - /** * Utility to print info/debug messages via InfoStream. * diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMergeReplay.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMergeReplay.java new file mode 100644 index 0000000000..c2bd07b87f --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMergeReplay.java @@ -0,0 +1,147 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix; +import java.io.IOException; +import java.util.List; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.KnnVectorValues; +import org.apache.lucene.tests.util.LuceneTestCase; + +public class TestAcceleratedHNSWMergeReplay extends LuceneTestCase { + + public void testUnderflowClosesBuilderBeforeBuild() { + TrackingBuilder builder = new TrackingBuilder(); + FloatVectorValues values = FloatVectorValues.fromFloats(List.of(new float[] {1f}), 1); + + IOException failure = + expectThrows( + IOException.class, + () -> Lucene99AcceleratedHNSWVectorsWriter.buildMergedDataset(values, 2, builder)); + + assertEquals( + "Merged vector count changed between passes: expected 2, observed 1", failure.getMessage()); + assertEquals(1, builder.addCalls); + assertEquals(0, builder.buildCalls); + assertEquals(1, builder.closeCalls); + } + + public void testOverflowClosesBuilderBeforeBuild() { + TrackingBuilder builder = new TrackingBuilder(); + FloatVectorValues values = + FloatVectorValues.fromFloats(List.of(new float[] {1f}, new float[] {2f}), 1); + + IOException failure = + expectThrows( + IOException.class, + () -> Lucene99AcceleratedHNSWVectorsWriter.buildMergedDataset(values, 1, builder)); + + assertEquals( + "Merged vector count changed between passes: expected 1, observed at least 2", + failure.getMessage()); + assertEquals(1, builder.addCalls); + assertEquals(0, builder.buildCalls); + assertEquals(1, builder.closeCalls); + } + + public void testReplayFailureRemainsPrimaryWhenBuilderCloseFails() { + IOException replayFailure = new IOException("vector read failed"); + IllegalStateException closeFailure = new IllegalStateException("builder close failed"); + TrackingBuilder builder = new TrackingBuilder(closeFailure); + FloatVectorValues values = failingValues(replayFailure); + + IOException thrown = + expectThrows( + IOException.class, + () -> Lucene99AcceleratedHNSWVectorsWriter.buildMergedDataset(values, 1, builder)); + + assertSame(replayFailure, thrown); + assertArrayEquals(new Throwable[] {closeFailure}, thrown.getSuppressed()); + assertEquals(0, builder.addCalls); + assertEquals(0, builder.buildCalls); + assertEquals(1, builder.closeCalls); + } + + private static FloatVectorValues failingValues(IOException failure) { + FloatVectorValues delegate = FloatVectorValues.fromFloats(List.of(new float[] {1f}), 1); + return new FloatVectorValues() { + @Override + public int dimension() { + return 1; + } + + @Override + public int size() { + return 1; + } + + @Override + public float[] vectorValue(int ord) throws IOException { + throw failure; + } + + @Override + public FloatVectorValues copy() { + return this; + } + + @Override + public KnnVectorValues.DocIndexIterator iterator() { + return delegate.iterator(); + } + }; + } + + private static final class TrackingBuilder implements CuVSMatrix.Builder { + private final RuntimeException closeFailure; + private int addCalls; + private int buildCalls; + private int closeCalls; + + private TrackingBuilder() { + this(null); + } + + private TrackingBuilder(RuntimeException closeFailure) { + this.closeFailure = closeFailure; + } + + @Override + public void addVector(float[] vector) { + addCalls++; + } + + @Override + public void addVector(byte[] vector) { + throw new AssertionError("unexpected byte vector"); + } + + @Override + public void addVector(int[] vector) { + throw new AssertionError("unexpected int vector"); + } + + @Override + public void addVector(short[] vector) { + throw new AssertionError("unexpected short vector"); + } + + @Override + public CuVSHostMatrix build() { + buildCalls++; + throw new AssertionError("build must not be called"); + } + + @Override + public void close() { + closeCalls++; + if (closeFailure != null) { + throw closeFailure; + } + } + } +} From 3897c69de0ae0fbd36219854b463538a994a8636 Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Sun, 20 Sep 2026 02:49:17 +0000 Subject: [PATCH 09/21] Restore fixed HNSW serialization waves --- .../cuvs/lucene/AcceleratedHNSWUtils.java | 31 +---- .../TestWriterThreadsGraphSerialization.java | 130 +++++++++++------- 2 files changed, 91 insertions(+), 70 deletions(-) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 46ed0a0d8e..6d45845e39 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -234,12 +234,9 @@ public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorInde int numLevels = graph.numLevels(); int[][] offsets = new int[numLevels][]; - // Level 0 holds all nodes and dominates serialization cost. Each node's delta/VInt block is - // independent, so encode level 0 in parallel and concatenate the per-thread buffers serially in - // node order, in memory-bounded waves. Higher levels are tiny and stay serial. The on-disk - // bytes - // are identical to the fully-serial path (blocks in node order, offsets = per-node byte - // lengths). + // Each level-0 node's delta/VInt block is independent. Encode those blocks in parallel in fixed + // waves, then concatenate thread buffers in node order. Higher levels are much smaller and stay + // serial. Both paths write the same blocks and offsets in the same order. // graph.maxConn() scans every layer-0 adjacency row (O(graph size)); compute it once here // rather than per level/per task below. int maxConn = graph.maxConn(); @@ -265,11 +262,8 @@ public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorInde /** Node count below which parallel level-0 serialization is not worth the overhead. */ static final int PARALLEL_MIN_NODES = 1 << 16; - /** Maximum encoded payload held by one parallel wave before it is copied to the index output. */ - static final long MAX_PARALLEL_ENCODE_BYTES = 64L << 20; - - /** Maximum bytes written by one node: one VInt count and {@code maxConn} VInt deltas. */ - private static final int MAX_VINT_BYTES = 5; + /** Nodes per wave, bounding the number of nodes buffered independently of dataset size. */ + static final int SERIALIZATION_WAVE_NODES = 1 << 20; /** Serially encodes a level's nodes into {@code out}, recording per-node byte lengths. */ private static void writeLevelSerial( @@ -291,7 +285,7 @@ private static void writeLevelSerial( } /** - * Encodes level 0 in parallel: within memory-bounded waves, threads encode contiguous node + * Encodes level 0 in parallel: within fixed-size waves, threads encode contiguous node * sub-ranges into per-thread buffers, which are then concatenated to {@code out} in node order * (identical layout to the serial path). */ @@ -312,9 +306,8 @@ private static void writeLevel0Parallel( try { TaskExecutor executor = new TaskExecutor(pool); int n = nodes.length; - int waveNodes = nodesPerSerializationWave(maxConn); for (int waveStart = 0; waveStart < n; ) { - int waveEnd = (int) Math.min(n, (long) waveStart + waveNodes); + int waveEnd = (int) Math.min(n, (long) waveStart + SERIALIZATION_WAVE_NODES); int perThread = (waveEnd - waveStart + numThreads - 1) / numThreads; ByteBuffersDataOutput[] buffers = new ByteBuffersDataOutput[numThreads]; @@ -353,16 +346,6 @@ private static void writeLevel0Parallel( } } - /** - * Sizes serialization waves from an upper bound on encoded bytes rather than only node count, so - * increasing graph degree cannot create an unbounded transient heap allocation. - */ - static int nodesPerSerializationWave(int maxConn) { - long maxBytesPerNode = (Math.max(0L, maxConn) + 1L) * MAX_VINT_BYTES; - return (int) - Math.max(1L, Math.min(Integer.MAX_VALUE, MAX_PARALLEL_ENCODE_BYTES / maxBytesPerNode)); - } - /** * Sorts, delta-encodes and de-duplicates a node's neighbors and writes the block (VInt size + VInt * deltas) to {@code out}. Shared by the serial and parallel paths so encoding is identical. diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java index 5ddb59cf88..0da7382fdc 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java @@ -19,64 +19,61 @@ import org.apache.lucene.store.IndexInput; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.hnsw.NeighborArray; import org.junit.Test; -/** Verifies parallel level-zero graph serialization is byte-identical and memory-bounded. */ +/** Verifies parallel level-zero graph serialization is byte-identical to serial serialization. */ public class TestWriterThreadsGraphSerialization extends LuceneTestCase { private static final int NUM_NODES = AcceleratedHNSWUtils.PARALLEL_MIN_NODES + 1000; private static final int DEGREE = 12; - private static final int REPORTED_MAX_CONN = 512; private static final int NUM_THREADS = 4; @Test - public void parallelSerializationMatchesSerialAcrossWaves() throws Exception { + public void parallelSerializationMatchesSerial() throws Exception { int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(2)); try (CuVSMatrix matrix = new ArrayMatrix(adjacency); Directory dir = new ByteBuffersDirectory()) { GPUBuiltHnswGraph serialGraph = newSingleLayerGraph(matrix); GPUBuiltHnswGraph parallelGraph = newSingleLayerGraph(matrix); + assertSerialAndParallelMatch(serialGraph, parallelGraph, dir); + } + } - int[][] serialOffsets; - try (IndexOutput out = dir.createOutput("serial", IOContext.DEFAULT)) { - serialOffsets = AcceleratedHNSWUtils.writeGraph(serialGraph, out, 1); - } - int[][] parallelOffsets; - try (IndexOutput out = dir.createOutput("parallel", IOContext.DEFAULT)) { - parallelOffsets = AcceleratedHNSWUtils.writeGraph(parallelGraph, out, NUM_THREADS); - } - - assertEquals(serialOffsets.length, parallelOffsets.length); - for (int level = 0; level < serialOffsets.length; level++) { - assertArrayEquals(serialOffsets[level], parallelOffsets[level]); - } - assertArrayEquals(readAllBytes(dir, "serial"), readAllBytes(dir, "parallel")); + @Test + public void parallelSerializationMatchesSerialAcrossFixedWaveBoundary() throws Exception { + int numNodes = AcceleratedHNSWUtils.SERIALIZATION_WAVE_NODES + 1; + assertTrue(numNodes > AcceleratedHNSWUtils.SERIALIZATION_WAVE_NODES); - assertTrue( - "test must cross a serialization-wave boundary", - AcceleratedHNSWUtils.nodesPerSerializationWave(REPORTED_MAX_CONN) < NUM_NODES); + try (Directory dir = new ByteBuffersDirectory()) { + assertSerialAndParallelMatch(new LazyEmptyGraph(numNodes), new LazyEmptyGraph(numNodes), dir); } } - @Test - public void serializationWaveHonorsEncodedByteBudget() { - for (int degree : new int[] {1, 32, 88, 152, 512}) { - int nodes = AcceleratedHNSWUtils.nodesPerSerializationWave(degree); - long maximumEncodedBytes = (long) nodes * (degree + 1L) * 5L; - assertTrue(maximumEncodedBytes <= AcceleratedHNSWUtils.MAX_PARALLEL_ENCODE_BYTES); - assertTrue(nodes > 0); + private static void assertSerialAndParallelMatch( + GPUBuiltHnswGraph serialGraph, GPUBuiltHnswGraph parallelGraph, Directory dir) + throws Exception { + int[][] serialOffsets; + try (IndexOutput out = dir.createOutput("serial", IOContext.DEFAULT)) { + serialOffsets = AcceleratedHNSWUtils.writeGraph(serialGraph, out, 1); + } + int[][] parallelOffsets; + try (IndexOutput out = dir.createOutput("parallel", IOContext.DEFAULT)) { + parallelOffsets = AcceleratedHNSWUtils.writeGraph(parallelGraph, out, NUM_THREADS); } + + assertEquals(serialOffsets.length, parallelOffsets.length); + for (int level = 0; level < serialOffsets.length; level++) { + assertArrayEquals(serialOffsets[level], parallelOffsets[level]); + } + assertArrayEquals(readAllBytes(dir, "serial"), readAllBytes(dir, "parallel")); } private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency) throws IOException { - return new ReportedMaxConnGraph( - NUM_NODES, - /* dimensions= */ 4, - Arrays.asList((int[]) null), - List.of(layer0Adjacency), - REPORTED_MAX_CONN); + return new GPUBuiltHnswGraph( + NUM_NODES, /* dimensions= */ 4, Arrays.asList((int[]) null), List.of(layer0Adjacency), 1); } private static byte[] readAllBytes(Directory dir, String name) throws Exception { @@ -97,24 +94,65 @@ private static int[][] randomAdjacency(int numNodes, int degree, Random random) return adjacency; } - /** Inflates maxConn only to force this modest test graph through multiple bounded waves. */ - private static final class ReportedMaxConnGraph extends GPUBuiltHnswGraph { - private final int reportedMaxConn; + /** Supplies an empty graph lazily so the fixed wave boundary can be tested with little heap. */ + private static final class LazyEmptyGraph extends GPUBuiltHnswGraph { + private final int graphSize; + + LazyEmptyGraph(int graphSize) throws IOException { + super( + 0, + /* dimensions= */ 4, + Arrays.asList((int[]) null), + List.of(new ArrayMatrix(new int[0][])), + 1); + this.graphSize = graphSize; + } - ReportedMaxConnGraph( - int size, - int dimensions, - List layerNodes, - List layerAdjacencies, - int reportedMaxConn) - throws IOException { - super(size, dimensions, layerNodes, layerAdjacencies, 1); - this.reportedMaxConn = reportedMaxConn; + @Override + public int size() { + return graphSize; } @Override public int maxConn() { - return reportedMaxConn; + return 0; + } + + @Override + public NodesIterator getNodesOnLevel(int level) { + return new RangeNodesIterator(level == 0 ? graphSize : 0); + } + + @Override + public NeighborArray getNeighbors(int level, int node) { + return null; + } + } + + private static final class RangeNodesIterator extends GPUBuiltHnswGraph.NodesIterator { + private int current = -1; + + RangeNodesIterator(int size) { + super(size); + } + + @Override + public boolean hasNext() { + return current + 1 < size; + } + + @Override + public int nextInt() { + return ++current; + } + + @Override + public int consume(int[] dest) { + int count = Math.min(dest.length, size - (current + 1)); + for (int i = 0; i < count; i++) { + dest[i] = ++current; + } + return count; } } From 7a58fb1fbb7f2134dd9d5811a1fb6b5b11370bbf Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Tue, 22 Sep 2026 19:50:06 +0000 Subject: [PATCH 10/21] Clarify accelerated HNSW matrix documentation --- .../nvidia/cuvs/lucene/AcceleratedHNSWUtils.java | 14 ++++++-------- .../main/java/com/nvidia/cuvs/lucene/Utils.java | 5 ++++- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 09325bc794..5f0b08b332 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -74,11 +74,9 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens } /** - * Creates a multi-layer HNSW graph with dynamic number of layers. - * M = ceil(cagraGraphDegree / 2), where cagraGraphDegree is the CAGRA adjacency list's degree - * (its column count). Ceil is used to accommodate odd graph degrees. - * Each layer contains 1/M nodes from the previous layer - * Creates layers until the highest layer has ≤ M nodes + * Creates up to {@code hnswLayers} total layers. Layer 0 uses the full CAGRA graph. Each upper + * layer samples {@code max(2, floor(previousLayerSize / M))} nodes, where {@code M} is {@code + * ceil(layer-0 graph degree / 2)}. */ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, @@ -184,6 +182,8 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( QuantizationType quantization) throws Throwable { int size = Math.toIntExact(vectorDataset.size()); + // Matrix columns are the stored width: binary vectors are bit-packed, while scalar and float + // vectors store one value per dimension. int columns = Math.toIntExact(vectorDataset.columns()); List vectors = new AbstractList<>() { @@ -216,9 +216,7 @@ public int size() { quantization); } - /** - * Builds a CAGRA graph for a subset of binary quantized vectors - */ + /** Builds a CAGRA graph for a selected vector subset. */ private static CuVSMatrix buildCagraGraphForSubset( Object vectors, int[] selectedNodes, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java index 97bc646225..5547db2bb4 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -95,7 +95,10 @@ static CuVSHostMatrix createHostFloatMatrix(List data, int dimensions) } } - /** Builds a host-memory CuVSMatrix from a list of byte vectors. */ + /** + * Builds a host-memory CuVSMatrix from byte vectors without first materializing the list as an + * intermediate {@code byte[][]}. + */ static CuVSHostMatrix createHostByteMatrix(List data, int bytesPerVector) { try (CuVSMatrix.Builder builder = CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE)) { From 5b6b22c83d933b5105ff012bad2153ef1053eb90 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Tue, 22 Sep 2026 21:48:50 +0000 Subject: [PATCH 11/21] Harden and document accelerated HNSW graph copying --- ...vidia-cuvs-lucene-acceleratedhnswparams.md | 66 ++++++++-------- ...nvidia-cuvs-lucene-acceleratedhnswutils.md | 28 ++++--- ...om-nvidia-cuvs-lucene-gpubuilthnswgraph.md | 44 ++++++++--- ...ne-lucene99acceleratedhnswvectorswriter.md | 10 +-- ...leratedhnswbinaryquantizedvectorswriter.md | 10 +-- ...leratedhnswscalarquantizedvectorswriter.md | 10 +-- ...lucene-api-com-nvidia-cuvs-lucene-utils.md | 15 ++-- fern/pages/user_guide/lucene.md | 4 +- .../cuvs/lucene/AcceleratedHNSWParams.java | 11 ++- .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 42 ++++++++--- ...TestWriterThreadsGraphMaterialization.java | 75 ++++++++++++++++++- 11 files changed, 217 insertions(+), 98 deletions(-) diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md index 27ffe5cb39..74661e14a8 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md @@ -18,13 +18,14 @@ public class AcceleratedHNSWParams public int getWriterThreads() ``` -Get the cuVS writer threads parameter +Get the maximum thread count for cuVS writes and accelerated-HNSW graph materialization and +serialization. **Returns** -cuVS writer threads parameter +maximum writer and graph-processing thread count -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:149`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:151`_ ### getIntermediateGraphDegree @@ -38,7 +39,7 @@ Get the intermediate graph degree the graph degree parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:158`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:160`_ ### getGraphdegree @@ -52,7 +53,7 @@ Get the graph degree the graph degree parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:167`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:169`_ ### getHnswLayers @@ -66,7 +67,7 @@ Get the number of HNSW layers the number of HNSW layers -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:176`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:178`_ ### getMaxConn @@ -80,7 +81,7 @@ Get the max connection parameter the max connection parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:185`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:187`_ ### getBeamWidth @@ -94,7 +95,7 @@ Get the beam width parameter the beam width parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:194`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:196`_ ### getCagraGraphBuildAlgo @@ -108,7 +109,7 @@ Get the CAGRA graph build algorithm the CAGRA graph build algorithm -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:203`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:205`_ ### getCuVSIvfPqParams @@ -122,7 +123,7 @@ Get the instance of `CuVSIvfPqParams` the instance of `CuVSIvfPqParams` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:212`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:214`_ ### getNumMergeWorkers @@ -136,7 +137,7 @@ Get the number of merge workers set to be used in the fallback mechanism the number of merge workers -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:221`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:223`_ ### getMergeExec @@ -150,7 +151,7 @@ Get the instance of the `ExecutorService` to be used in the fallback mechanism the instance of the `ExecutorService` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:230`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:232`_ ### getStrategy @@ -167,7 +168,7 @@ When CUSTOM is chosen, the build algorithm and its parameters (either defaults o get the chosen `Strategy` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:242`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:244`_ ### getCuvsDistanceType @@ -181,7 +182,7 @@ Get the cuvs distance type the distance type -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:251`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:253`_ ### getNNDescentNumIterations @@ -195,7 +196,7 @@ get the number of Iterations to run if building with NN_DESCENT the number of iterations for NN_DESCENT -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:260`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:262`_ ### getHnswHeuristicType @@ -210,7 +211,7 @@ beamWidth. Only consulted under the `Strategy#HEURISTIC` strategy. the `HnswHeuristicType` to hand to cuVS -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:270`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:272`_ ### withWriterThreads @@ -218,7 +219,8 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWP public Builder withWriterThreads(int writerThreads) ``` -Set the number of cuVS writer threads while building the index +Set the maximum number of threads used for cuVS writes and accelerated-HNSW graph +materialization and serialization. Valid range - Minimum: \{@value MIN_WRITER_THREADS\}, Maximum: \{@value MAX_WRITER_THREADS\} Default value - \{@value DEFAULT_WRITER_THREADS\} @@ -232,7 +234,7 @@ Default value - \{@value DEFAULT_WRITER_THREADS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:335`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:338`_ ### withIntermediateGraphDegree @@ -254,7 +256,7 @@ Default value - \{@value DEFAULT_INT_GRAPH_DEGREE\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:348`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:351`_ ### withGraphDegree @@ -276,7 +278,7 @@ Default value - \{@value DEFAULT_GRAPH_DEGREE\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:361`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:364`_ ### withHNSWLayer @@ -298,7 +300,7 @@ Default value - \{@value DEFAULT_HNSW_LAYERS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:374`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:377`_ ### withMaxConn @@ -320,7 +322,7 @@ Default value - \{@value DEFAULT_MAX_CONN\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:387`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:390`_ ### withBeamWidth @@ -342,7 +344,7 @@ Default value - \{@value DEFAULT_BEAM_WIDTH\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:400`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:403`_ ### withCagraGraphBuildAlgo @@ -363,7 +365,7 @@ Default value - NN_DESCENT instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:412`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:415`_ ### withCuVSIvfPqParams @@ -383,7 +385,7 @@ Set the instance of `CuVSIvfPqParams` instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:423`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:426`_ ### withNumMergeWorkers @@ -404,7 +406,7 @@ Default value - \{@value DEFAULT_NUM_MERGE_WORKERS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:435`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:438`_ ### withMergeExecutorService @@ -425,7 +427,7 @@ Default value an instance with one thread instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:447`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:450`_ ### withStrategy @@ -451,7 +453,7 @@ Default value - HEURISTIC instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:464`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:467`_ ### withCuvsDistanceType @@ -471,7 +473,7 @@ Set the CuvsDistanceType instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:475`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:478`_ ### withNNDescentNumIterations @@ -494,7 +496,7 @@ Default value - \{@value DEFAULT_NN_DESCENT_NUM_ITERATIONS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:489`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:492`_ ### withHnswHeuristicType @@ -518,7 +520,7 @@ the equivalent HNSW graph (graph degree = 2 * maxConn). instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:504`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:507`_ ### build @@ -532,6 +534,6 @@ Create an instance of `AcceleratedHNSWParams` instance of `AcceleratedHNSWParams` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:600`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:603`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:17`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md index b3bd116a4c..1af0d77d1f 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md @@ -21,7 +21,7 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens Creates a dummy HNSW graph for a single vector. The graph will have 1 level with 1 node and no neighbors. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:55`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:61`_ ### createMultiLayerHnswGraph @@ -29,13 +29,11 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWU public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, int size, int dimensions, CuVSMatrix adjacencyListMatrix, List vectors, int hnswLayers, CagraIndexParams params, QuantizationType quantization) throws Throwable ``` -Creates a multi-layer HNSW graph with dynamic number of layers. -M = ceil(cagraGraphDegree / 2), where cagraGraphDegree is the CAGRA adjacency list's degree -(its column count). Ceil is used to accommodate odd graph degrees. -Each layer contains 1/M nodes from the previous layer -Creates layers until the highest layer has ≤ M nodes +Creates up to `hnswLayers` total layers. Layer 0 uses the full CAGRA graph. Each upper +layer samples `max(2, floor(previousLayerSize / M))` nodes, where `M` is \{@code +ceil(layer-0 graph degree / 2)\}. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:82`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:86`_ ### createMultiLayerHnswGraph @@ -46,7 +44,7 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, int dim Creates a multi-layer HNSW graph from a native matrix without copying the complete dataset to the Java heap. The list view copies only rows selected for an upper layer. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:176`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:203`_ ### writeGraph @@ -73,7 +71,7 @@ a 2D array of offsets | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:288`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:337`_ ### writeMeta @@ -102,7 +100,7 @@ Writes the meta information for the index. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:353`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:501`_ ### printInfoStream @@ -118,7 +116,7 @@ A utility method to print info/debugging messages using InfoStream. | --- | --- | | `msg` | the debugging message to print | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:435`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:583`_ ### writeEmpty @@ -140,7 +138,7 @@ Writes an empty meta information for the field. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:447`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:595`_ ### quantizeFloatVectorsToBinary @@ -163,7 +161,7 @@ Bits are packed: 8 dimensions per byte. A list of byte binary representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:460`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:608`_ ### quantizeFloatVectorsToScalar @@ -183,6 +181,6 @@ Scalar quantization. A list of byte scalar representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:502`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:650`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:32`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:38`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md index f8a3aee9dc..6425b37388 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md @@ -31,7 +31,27 @@ Multi-layer constructor that supports arbitrary number of layers. | `layerNodes` | the nodes on the layer | | `layerAdjacencies` | adjacency list | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:41`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:55`_ + +### GPUBuiltHnswGraph + +```java +public GPUBuiltHnswGraph( int size, int dimensions, List layerNodes, List layerAdjacencies, int numThreads) throws IOException +``` + +Builds a graph while materializing adjacency rows with the requested number of threads. + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:61`_ + +### fitsParallelGraphCopyBudget + +```java +static boolean fitsParallelGraphCopyBudget(long rows, long columns) +``` + +Returns whether an INT32 adjacency can be copied without exceeding the native-host budget. + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:169`_ ### getNodesOnLevel @@ -41,7 +61,7 @@ public NodesIterator getNodesOnLevel(int level) Get all nodes on a given level as node 0th ordinals. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:89`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:253`_ ### getNeighbors @@ -62,7 +82,7 @@ Get the neighbors for the node and the level it resides. an instance of NeighborArray -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:107`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:271`_ ### seek @@ -72,7 +92,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Move the pointer to exactly the given level's target. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:132`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:296`_ ### nextNeighbor @@ -82,7 +102,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Iterates over the neighbor list. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:142`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:306`_ ### entryNode @@ -92,7 +112,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns graph's entry point on the top level. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:173`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:337`_ ### maxConn @@ -102,7 +122,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap returns M, the maximum number of connections for a node. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:192`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:356`_ ### neighborCount @@ -112,7 +132,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns the neighbor count. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:207`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:371`_ ### size @@ -122,7 +142,7 @@ public int size() Returns the number of nodes in the graph. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:282`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:446`_ ### numLevels @@ -136,7 +156,7 @@ Returns the number of levels in the HNSW graph. the number of levels -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:291`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:455`_ ### dimensions @@ -150,6 +170,6 @@ Gets the vector dimension. the vector dimension -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:300`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:464`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:21`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:29`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md index 7f5698a6a3..a7a101bc8d 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:222`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:224`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:364`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:367`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:373`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:376`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:393`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:396`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:403`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:406`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:55`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md index 83a7741bed..67a30a816d 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:213`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:215`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:300`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:303`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:331`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:334`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:351`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:354`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:360`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:363`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:57`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md index 66d4678107..e90827ba09 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:239`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:241`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:324`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:327`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:355`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:358`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:375`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:378`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:384`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:387`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:56`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md index 5717013e42..03e59814ff 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md @@ -102,9 +102,10 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:87`_ static CuVSHostMatrix createHostByteMatrix(List data, int bytesPerVector) ``` -Builds a host-memory CuVSMatrix from a list of byte vectors. +Builds a host-memory CuVSMatrix from byte vectors without first materializing the list as an +intermediate `byte[][]`. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:98`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:101`_ ### createHostByteMatrixFromArray @@ -114,7 +115,7 @@ static CuVSHostMatrix createHostByteMatrixFromArray(byte[][] data, int bytesPerV Builds a host-memory CuVSMatrix from a 2D byte array. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:109`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:112`_ ### nanosToMillis @@ -134,7 +135,7 @@ A utility method to convert nanoseconds to milliseconds. milliseconds -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:125`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:128`_ ### cuVSResourcesOrNull @@ -148,7 +149,7 @@ Creates an instance of CuVSResources. an instance of CuVSResources -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:134`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:137`_ ### handleThrowableWithIgnore @@ -171,7 +172,7 @@ A utility method that conditionally ignores certain throwable objects | --- | --- | | `IOException` | | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:162`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:165`_ ### info @@ -189,6 +190,6 @@ Utility to print info/debug messages via InfoStream. | `component` | the name of the index writer | | `msg` | the log message to push via the InfoStream | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:176`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:179`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:22`_ diff --git a/fern/pages/user_guide/lucene.md b/fern/pages/user_guide/lucene.md index 34aa9402d4..2f3e49cd20 100644 --- a/fern/pages/user_guide/lucene.md +++ b/fern/pages/user_guide/lucene.md @@ -231,7 +231,9 @@ Both `AcceleratedHNSWParams` and `GPUSearchParams` default to a `HEURISTIC` stra - For the accelerated HNSW codecs, set `maxConn` and `beamWidth`, the HNSW parameters you would tune on the CPU. cuVS derives graph degrees and the build algorithm from them. These two values also configure the CPU fallback writer, so one setting covers both paths. - For the GPU search codec, set `buildQuality`. Higher values spend more build time for a higher-quality graph. This codec also passes `graphDegree` into the heuristic, so leave it at its default unless you intend to cap the graph. -Increasing `writerThreads` raises index build concurrency. The accelerated HNSW codecs default to a single writer thread, while the GPU search codec defaults to 32. +`writerThreads` controls native cuVS writer concurrency for both parameter types. For the accelerated HNSW codecs, it also sets the maximum number of CPU threads used to materialize and serialize the finished graph; the default of one keeps this post-build processing serial. The GPU search codec continues to use the setting only for native writer concurrency and defaults to 32. + +With more than one writer thread, parallel adjacency materialization is eligible for layers with at least 65,536 nodes. A device adjacency is copied once to native host memory when that copy is at most 4 GiB; larger device adjacencies retain serial row access. Level-zero serialization is independently eligible at 65,536 nodes and runs in fixed waves of at most 1,048,576 nodes; upper levels remain serial. Switching either class to the `CUSTOM` strategy exposes the underlying CAGRA parameters directly, including `graphDegree`, `intermediateGraphDegree`, the graph build algorithm, and its parameters. Use `CUSTOM` only when you have measurements that justify specific values; the defaults derived by cuVS are a better starting point. For background on the parameters themselves, see the [CAGRA indexing guide](/user-guide/api-guides/indexing-guide/cagra) and the [tuning guide](/getting-started/introduction/tuning-indexes). diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index a5f164b70b..f63191c091 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -95,7 +95,8 @@ public static enum Strategy { /** * Constructs an instance of {@link AcceleratedHNSWParams} with specific parameter values. * - * @param writerThreads Number of cuVS writer threads to use. + * @param writerThreads Maximum number of threads to use for cuVS writes and accelerated-HNSW + * graph materialization and serialization. * @param intermediateGraphDegree The intermediate graph degree while building the CAGRA index. * @param graphdegree The graph degree to use while building the CAGRA index. * @param hnswLayers The number of HNSW layers to build in the HNSW index. @@ -143,9 +144,10 @@ private AcceleratedHNSWParams( } /** - * Get the cuVS writer threads parameter + * Get the maximum thread count for cuVS writes and accelerated-HNSW graph materialization and + * serialization. * - * @return cuVS writer threads parameter + * @return maximum writer and graph-processing thread count */ public int getWriterThreads() { return writerThreads; @@ -326,7 +328,8 @@ public static class Builder { private HnswHeuristicType hnswHeuristicType = DEFAULT_HNSW_HEURISTIC_TYPE; /** - * Set the number of cuVS writer threads while building the index + * Set the maximum number of threads used for cuVS writes and accelerated-HNSW graph + * materialization and serialization. * Valid range - Minimum: {@value MIN_WRITER_THREADS}, Maximum: {@value MAX_WRITER_THREADS} * Default value - {@value DEFAULT_WRITER_THREADS} * diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index 9e61aed5c2..be9a8c5846 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -16,6 +16,7 @@ import java.util.concurrent.Callable; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.function.Supplier; import org.apache.lucene.search.TaskExecutor; import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.NeighborArray; @@ -149,20 +150,14 @@ private static NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size, fillNeighborRange(adjacency, neighbors, 0, size); return neighbors; } - CuVSMatrix source = adjacency; - CuVSHostMatrix hostCopy = null; if (adjacency instanceof CuVSDeviceMatrix deviceAdjacency) { - hostCopy = deviceAdjacency.toHost(); - source = hostCopy; - } - try { - fillNeighborArrayParallel(source, neighbors, size, numThreads); - return neighbors; - } finally { - if (hostCopy != null) { - hostCopy.close(); + try (CuVSHostMatrix hostCopy = copyToHost(deviceAdjacency)) { + fillNeighborArrayParallel(hostCopy, neighbors, size, numThreads); } + return neighbors; } + fillNeighborArrayParallel(adjacency, neighbors, size, numThreads); + return neighbors; } private static NeighborArray[] fillNeighborArraySerial(CuVSMatrix adjacency, int size) { @@ -182,6 +177,31 @@ static boolean fitsParallelGraphCopyBudget(long rows, long columns) { return rows <= MAX_PARALLEL_GRAPH_COPY_BYTES / Integer.BYTES / columns; } + private static CuVSHostMatrix copyToHost(CuVSDeviceMatrix source) { + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(source.size(), source.columns(), source.dataType())) { + return copyToHost(source, builder::build); + } + } + + static CuVSHostMatrix copyToHost( + CuVSDeviceMatrix source, Supplier hostCopyFactory) { + CuVSHostMatrix hostCopy = hostCopyFactory.get(); + try { + source.toHost(hostCopy); + return hostCopy; + } catch (RuntimeException | Error failure) { + try { + hostCopy.close(); + } catch (RuntimeException | Error closeFailure) { + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } + } + throw failure; + } + } + /** * Materializes disjoint node ranges concurrently. Each thread writes its own slots of {@code * neighbors} and its own {@link NeighborArray} instances, so no synchronization is needed; {@code diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java index 03cb017c18..ad3366dd14 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java @@ -16,6 +16,7 @@ import java.util.Arrays; import java.util.List; import java.util.Random; +import java.util.concurrent.atomic.AtomicInteger; import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; @@ -59,6 +60,49 @@ public void oversizedDeviceAdjacencyUsesSerialFallback() throws Exception { } } + @Test + public void failedDeviceCopyClosesHostAllocationAndSuppressesCloseFailure() { + RuntimeException copyFailure = new RuntimeException("copy failed"); + RuntimeException closeFailure = new RuntimeException("close failed"); + AtomicInteger hostCloseCount = new AtomicInteger(); + CuVSDeviceMatrix source = + new ArrayDeviceMatrix(new int[][] {{0}}, 1) { + @Override + public void toHost(CuVSHostMatrix target) { + throw copyFailure; + } + }; + CuVSHostMatrix hostCopy = new TrackingHostMatrix(hostCloseCount, closeFailure); + + RuntimeException thrown = + assertThrows( + RuntimeException.class, () -> GPUBuiltHnswGraph.copyToHost(source, () -> hostCopy)); + + assertSame(copyFailure, thrown); + assertEquals(1, hostCloseCount.get()); + assertArrayEquals(new Throwable[] {closeFailure}, thrown.getSuppressed()); + } + + @Test + public void successfulDeviceCopyTransfersHostOwnershipToCaller() { + AtomicInteger hostCloseCount = new AtomicInteger(); + CuVSHostMatrix hostCopy = new TrackingHostMatrix(hostCloseCount, null); + CuVSDeviceMatrix source = + new ArrayDeviceMatrix(new int[][] {{0}}, 1) { + @Override + public void toHost(CuVSHostMatrix target) { + assertSame(hostCopy, target); + } + }; + + CuVSHostMatrix returned = GPUBuiltHnswGraph.copyToHost(source, () -> hostCopy); + + assertSame(hostCopy, returned); + assertEquals(0, hostCloseCount.get()); + returned.close(); + assertEquals(1, hostCloseCount.get()); + } + private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, int numThreads) throws IOException { return new GPUBuiltHnswGraph( @@ -170,7 +214,7 @@ public void close() {} } /** Reports an oversized device shape and fails if the guarded host-copy path is reached. */ - private static final class ArrayDeviceMatrix extends ArrayMatrix implements CuVSDeviceMatrix { + private static class ArrayDeviceMatrix extends ArrayMatrix implements CuVSDeviceMatrix { private final long reportedColumns; ArrayDeviceMatrix(int[][] rows, long reportedColumns) { @@ -183,12 +227,41 @@ public long columns() { return reportedColumns; } + @Override + public void toHost(CuVSHostMatrix target) { + throw new AssertionError("oversized device adjacency must not be copied to host"); + } + @Override public CuVSHostMatrix toHost() { throw new AssertionError("oversized device adjacency must not be copied to host"); } } + private static final class TrackingHostMatrix extends ArrayMatrix implements CuVSHostMatrix { + private final AtomicInteger closeCount; + private final RuntimeException closeFailure; + + TrackingHostMatrix(AtomicInteger closeCount, RuntimeException closeFailure) { + super(new int[][] {{0}}); + this.closeCount = closeCount; + this.closeFailure = closeFailure; + } + + @Override + public int get(int row, int column) { + return 0; + } + + @Override + public void close() { + closeCount.incrementAndGet(); + if (closeFailure != null) { + throw closeFailure; + } + } + } + private static final class ArrayRow implements RowView { private final int[] values; From ddd430e16706b4f16129bea97420fe052df0ead4 Mon Sep 17 00:00:00 2001 From: EC2 Default User Date: Tue, 22 Sep 2026 23:08:20 +0000 Subject: [PATCH 12/21] Refresh Lucene API documentation --- ...nvidia-cuvs-lucene-acceleratedhnswutils.md | 24 +++++++++---------- ...lucene-api-com-nvidia-cuvs-lucene-utils.md | 15 ++++++------ .../cuvs/lucene/AcceleratedHNSWUtils.java | 4 ++-- 3 files changed, 21 insertions(+), 22 deletions(-) diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md index b3bd116a4c..102ed72776 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md @@ -29,13 +29,11 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWU public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, int size, int dimensions, CuVSMatrix adjacencyListMatrix, List vectors, int hnswLayers, CagraIndexParams params, QuantizationType quantization) throws Throwable ``` -Creates a multi-layer HNSW graph with dynamic number of layers. -M = ceil(cagraGraphDegree / 2), where cagraGraphDegree is the CAGRA adjacency list's degree -(its column count). Ceil is used to accommodate odd graph degrees. -Each layer contains 1/M nodes from the previous layer -Creates layers until the highest layer has ≤ M nodes +Creates up to `hnswLayers` total layers. Layer 0 uses the full CAGRA graph. Each upper +layer samples `max(2, floor(previousLayerSize / M))` nodes. The value `M` is the +ceiling of half the layer-0 graph degree. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:82`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:80`_ ### createMultiLayerHnswGraph @@ -46,7 +44,7 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, int dim Creates a multi-layer HNSW graph from a native matrix without copying the complete dataset to the Java heap. The list view copies only rows selected for an upper layer. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:176`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:174`_ ### writeGraph @@ -73,7 +71,7 @@ a 2D array of offsets | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:288`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:286`_ ### writeMeta @@ -102,7 +100,7 @@ Writes the meta information for the index. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:353`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:351`_ ### printInfoStream @@ -118,7 +116,7 @@ A utility method to print info/debugging messages using InfoStream. | --- | --- | | `msg` | the debugging message to print | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:435`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:433`_ ### writeEmpty @@ -140,7 +138,7 @@ Writes an empty meta information for the field. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:447`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:445`_ ### quantizeFloatVectorsToBinary @@ -163,7 +161,7 @@ Bits are packed: 8 dimensions per byte. A list of byte binary representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:460`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:458`_ ### quantizeFloatVectorsToScalar @@ -183,6 +181,6 @@ Scalar quantization. A list of byte scalar representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:502`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:500`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:32`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md index 5717013e42..03e59814ff 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md @@ -102,9 +102,10 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:87`_ static CuVSHostMatrix createHostByteMatrix(List data, int bytesPerVector) ``` -Builds a host-memory CuVSMatrix from a list of byte vectors. +Builds a host-memory CuVSMatrix from byte vectors without first materializing the list as an +intermediate `byte[][]`. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:98`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:101`_ ### createHostByteMatrixFromArray @@ -114,7 +115,7 @@ static CuVSHostMatrix createHostByteMatrixFromArray(byte[][] data, int bytesPerV Builds a host-memory CuVSMatrix from a 2D byte array. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:109`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:112`_ ### nanosToMillis @@ -134,7 +135,7 @@ A utility method to convert nanoseconds to milliseconds. milliseconds -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:125`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:128`_ ### cuVSResourcesOrNull @@ -148,7 +149,7 @@ Creates an instance of CuVSResources. an instance of CuVSResources -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:134`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:137`_ ### handleThrowableWithIgnore @@ -171,7 +172,7 @@ A utility method that conditionally ignores certain throwable objects | --- | --- | | `IOException` | | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:162`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:165`_ ### info @@ -189,6 +190,6 @@ Utility to print info/debug messages via InfoStream. | `component` | the name of the index writer | | `msg` | the log message to push via the InfoStream | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:176`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:179`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:22`_ diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 5f0b08b332..52a28a0a6d 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -75,8 +75,8 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens /** * Creates up to {@code hnswLayers} total layers. Layer 0 uses the full CAGRA graph. Each upper - * layer samples {@code max(2, floor(previousLayerSize / M))} nodes, where {@code M} is {@code - * ceil(layer-0 graph degree / 2)}. + * layer samples {@code max(2, floor(previousLayerSize / M))} nodes. The value {@code M} is the + * ceiling of half the layer-0 graph degree. */ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, From 502d977e6b8a5de3f219d16568cd1d74b3bb4dd0 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Wed, 23 Sep 2026 16:44:07 +0000 Subject: [PATCH 13/21] Fix accelerated HNSW resource lifecycles --- .../java-api-com-nvidia-cuvs-cagraindex.md | 72 ++--- ...nvidia-cuvs-lucene-acceleratedhnswutils.md | 18 +- ...om-nvidia-cuvs-lucene-gpubuilthnswgraph.md | 20 +- ...ne-lucene99acceleratedhnswvectorswriter.md | 10 +- ...leratedhnswbinaryquantizedvectorswriter.md | 10 +- ...leratedhnswscalarquantizedvectorswriter.md | 10 +- ...lucene-api-com-nvidia-cuvs-lucene-utils.md | 20 +- .../main/java/com/nvidia/cuvs/CagraIndex.java | 8 +- .../cuvs/lucene/AcceleratedHNSWUtils.java | 249 +++++++++++------- .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 5 +- .../Lucene99AcceleratedHNSWVectorsWriter.java | 7 +- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 51 ++-- ...ratedHNSWScalarQuantizedVectorsWriter.java | 50 ++-- .../java/com/nvidia/cuvs/lucene/Utils.java | 66 +++++ .../TestAcceleratedHNSWDeletedDocuments.java | 5 + .../TestAcceleratedHNSWUpperLayers.java | 14 +- .../lucene/TestUtilsThrowableHandling.java | 66 +++++ 17 files changed, 451 insertions(+), 230 deletions(-) diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindex.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindex.md index b26f36e2d1..c89f50b924 100644 --- a/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindex.md +++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-cagraindex.md @@ -97,9 +97,10 @@ _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:106`_ @Override void close() throws Exception ``` -Invokes the native destroy_cagra_index to de-allocate the CAGRA index +Invokes the native destroy_cagra_index to de-allocate the CAGRA index. Also attempts to close +any dataset whose ownership transferred to this index during construction. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:136`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:137`_ ### search @@ -120,7 +121,7 @@ CAGRA index. an instance of `SearchResults` containing the results -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:147`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:148`_ ### makePaddedDataset @@ -132,7 +133,7 @@ Create an owning padded dataset by allocating padded storage and copying `dataset`. Prefer this when the source matrix is not already padded to CAGRA's required row stride (e.g. unaligned dimensions). -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:154`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:155`_ ### makePaddedDatasetView @@ -144,7 +145,7 @@ Create a caller-owned padded dataset view handle from a matrix that is already padded to CAGRA's required row stride. For unpadded matrices use `#makePaddedDataset(CuVSMatrix)`. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:161`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:162`_ ### makeStandardDatasetView @@ -154,7 +155,7 @@ StandardDatasetView makeStandardDatasetView(CuVSMatrix dataset) throws Throwable Create a caller-owned standard dataset view handle from a matrix. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:164`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:165`_ ### updateDataset @@ -166,7 +167,7 @@ Update this index with a caller-provided padded device dataset view and leave it search-ready in padded-device layout. The caller retains ownership of the underlying padded storage and must keep it alive while this index uses it. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:171`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:172`_ ### updateDataset @@ -177,7 +178,7 @@ void updateDataset(PaddedDataset dataset) throws Throwable Update this index with a caller-owned padded device dataset. The dataset must remain alive while this index uses it. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:177`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:178`_ ### getGraph @@ -191,7 +192,7 @@ Returns the CAGRA graph a `CuVSDeviceMatrix` encapsulating the native int (uint32_t) array used to represent the cagra graph -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:184`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:185`_ ### getGraphDegree @@ -207,7 +208,7 @@ build truncated it. the built graph degree (`graph().extent(1)`) -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:193`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:194`_ ### size @@ -221,7 +222,7 @@ Returns the number of vectors in this index. the number of rows of the indexed dataset -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:200`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:201`_ ### serialize @@ -238,7 +239,7 @@ for writing index bytes. | --- | --- | | `outputStream` | an instance of `OutputStream` to write the index bytes into | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:209`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:210`_ ### serialize @@ -256,7 +257,7 @@ for writing index bytes. | `outputStream` | an instance of `OutputStream` to write the index bytes into | | `bufferLength` | the length of buffer to use for writing bytes. Default value is 1024 | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:220`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:221`_ ### serialize @@ -274,7 +275,7 @@ for writing index bytes. | `outputStream` | an instance of `OutputStream` to write the index bytes into | | `tempFile` | an intermediate `Path` where CAGRA index is written temporarily | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:231`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:232`_ ### serialize @@ -293,7 +294,7 @@ and path to the intermediate temporary file. | `tempFile` | an intermediate `Path` where CAGRA index is written temporarily | | `bufferLength` | the length of buffer to use for writing bytes. Default value is 1024 | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:246`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:247`_ ### serializeToHNSW @@ -310,7 +311,7 @@ of `OutputStream` and path to the intermediate temporary file. | --- | --- | | `outputStream` | an instance of `OutputStream` to write the index bytes to | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:255`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:256`_ ### serializeToHNSW @@ -328,7 +329,7 @@ of `OutputStream` and path to the intermediate temporary file. | `outputStream` | an instance of `OutputStream` to write the index bytes to | | `bufferLength` | the length of buffer to use for writing bytes. Default value is 1024 | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:266`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:267`_ ### serializeToHNSW @@ -346,7 +347,7 @@ of `OutputStream` and path to the intermediate temporary file. | `outputStream` | an instance of `OutputStream` to write the index bytes to | | `tempFile` | an intermediate `Path` where CAGRA index is written temporarily | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:277`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:278`_ ### serializeToHNSW @@ -365,7 +366,7 @@ of `OutputStream` and path to the intermediate temporary file. | `tempFile` | an intermediate `Path` where CAGRA index is written temporarily | | `bufferLength` | the length of buffer to use for writing bytes. Default value is 1024 | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:292`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:293`_ ### getCuVSResources @@ -379,7 +380,7 @@ Gets an instance of `CuVSResources` an instance of `CuVSResources` -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:299`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:300`_ ### newBuilder @@ -401,7 +402,7 @@ Creates a new Builder with an instance of `CuVSResources`. | --- | --- | | `UnsupportedOperationException` | if the provider does not cuvs | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:307`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:308`_ ### merge @@ -427,7 +428,7 @@ A new merged CAGRA index | --- | --- | | `Throwable` | if an error occurs during the merge operation | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:319`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:320`_ ### merge @@ -454,7 +455,7 @@ A new merged CAGRA index | --- | --- | | `Throwable` | if an error occurs during the merge operation | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:331`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:332`_ ### merge @@ -490,7 +491,7 @@ A new merged CAGRA index | `IllegalArgumentException` | if `rowFilter` has a bit set beyond the last row, or if it is non-null but keeps no rows at all | | `Throwable` | if an error occurs during the merge operation | -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:355`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:356`_ ### isPaddedDataset @@ -516,7 +517,7 @@ copy it into padded storage it already occupies, and one that is not has to go t true when the rows are already padded the way CAGRA requires -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:383`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:384`_ ### from @@ -537,7 +538,7 @@ needed. an instance of this Builder -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:400`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:401`_ ### from @@ -560,7 +561,7 @@ serialized index. Keep `outDataset` alive while the built index is in use. an instance of this Builder -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:411`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:412`_ ### from @@ -571,7 +572,7 @@ Builder from(CuVSMatrix graph) Sets a CAGRA graph instance to re-create an index from a previously built graph. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:417`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:418`_ ### withDataset @@ -591,7 +592,7 @@ Sets the dataset vectors for building the `CagraIndex`. an instance of this Builder -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:425`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:426`_ ### withDataset @@ -601,6 +602,11 @@ Builder withDataset(CuVSMatrix dataset) Sets the dataset for building the `CagraIndex`. +The caller retains ownership until a build that uses this dataset returns successfully. +The returned index then owns the dataset, and the caller must leave it open until the index +is closed. If the build fails or uses another configured input source, ownership remains +with the caller. + **Parameters** | Name | Description | @@ -611,7 +617,7 @@ Sets the dataset for building the `CagraIndex`. an instance of this Builder -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:433`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:439`_ ### withBbqDataset @@ -629,7 +635,7 @@ stay open for as long as the index is in use. A dense dataset passed to `#withDataset(CuVSMatrix)` is owned by the index, as it is for a non-BBQ build, and is closed with it. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:446`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:452`_ ### withIndexParams @@ -650,7 +656,7 @@ Builder. An instance of this Builder. -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:455`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:461`_ ### build @@ -664,6 +670,6 @@ Builds and returns an instance of CagraIndex. an instance of CagraIndex -_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:462`_ +_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:468`_ _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java:26`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md index 102ed72776..aa65c2cf30 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md @@ -33,18 +33,18 @@ Creates up to `hnswLayers` total layers. Layer 0 uses the full CAGRA graph. Each layer samples `max(2, floor(previousLayerSize / M))` nodes. The value `M` is the ceiling of half the layer-0 graph degree. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:80`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:73`_ ### createMultiLayerHnswGraph ```java -static GPUBuiltHnswGraph createMultiLayerHnswGraph( FieldInfo fieldInfo, int dimensions, CuVSMatrix adjacencyListMatrix, CuVSMatrix vectorDataset, int hnswLayers, CagraIndexParams params, QuantizationType quantization) throws Throwable +static GPUBuiltHnswGraph createMultiLayerHnswGraph( int dimensions, CuVSMatrix adjacencyListMatrix, CuVSMatrix vectorDataset, int hnswLayers, CagraIndexParams params, QuantizationType quantization) throws Throwable ``` Creates a multi-layer HNSW graph from a native matrix without copying the complete dataset to the Java heap. The list view copies only rows selected for an upper layer. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:174`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:204`_ ### writeGraph @@ -71,7 +71,7 @@ a 2D array of offsets | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:286`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:333`_ ### writeMeta @@ -100,7 +100,7 @@ Writes the meta information for the index. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:351`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:398`_ ### printInfoStream @@ -116,7 +116,7 @@ A utility method to print info/debugging messages using InfoStream. | --- | --- | | `msg` | the debugging message to print | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:433`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:480`_ ### writeEmpty @@ -138,7 +138,7 @@ Writes an empty meta information for the field. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:445`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:492`_ ### quantizeFloatVectorsToBinary @@ -161,7 +161,7 @@ Bits are packed: 8 dimensions per byte. A list of byte binary representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:458`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:505`_ ### quantizeFloatVectorsToScalar @@ -181,6 +181,6 @@ Scalar quantization. A list of byte scalar representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:500`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:547`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:32`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md index f8a3aee9dc..fddfb07773 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md @@ -41,7 +41,7 @@ public NodesIterator getNodesOnLevel(int level) Get all nodes on a given level as node 0th ordinals. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:89`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:92`_ ### getNeighbors @@ -62,7 +62,7 @@ Get the neighbors for the node and the level it resides. an instance of NeighborArray -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:107`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:110`_ ### seek @@ -72,7 +72,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Move the pointer to exactly the given level's target. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:132`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:135`_ ### nextNeighbor @@ -82,7 +82,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Iterates over the neighbor list. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:142`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:145`_ ### entryNode @@ -92,7 +92,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns graph's entry point on the top level. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:173`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:176`_ ### maxConn @@ -102,7 +102,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap returns M, the maximum number of connections for a node. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:192`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:195`_ ### neighborCount @@ -112,7 +112,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns the neighbor count. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:207`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:210`_ ### size @@ -122,7 +122,7 @@ public int size() Returns the number of nodes in the graph. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:282`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:285`_ ### numLevels @@ -136,7 +136,7 @@ Returns the number of levels in the HNSW graph. the number of levels -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:291`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:294`_ ### dimensions @@ -150,6 +150,6 @@ Gets the vector dimension. the vector dimension -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:300`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:303`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:21`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md index 7f5698a6a3..150ef20df2 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:222`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:221`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:364`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:363`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:373`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:372`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:393`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:392`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:403`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:402`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:55`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md index 83a7741bed..b9ef604a4a 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:213`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:220`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:300`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:305`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:331`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:336`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:351`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:356`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:360`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:365`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:57`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md index 66d4678107..fc8c021fb0 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:239`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:247`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:324`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:330`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:355`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:361`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:375`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:381`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:384`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:390`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:56`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md index 03e59814ff..05fa15a980 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md @@ -45,7 +45,7 @@ never returns; always throws | --- | --- | | `IOException` | | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:39`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:40`_ ### createFloatMatrix @@ -70,7 +70,7 @@ without creating intermediate heap arrays. an instance of CuVSMatrix -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:59`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:60`_ ### createHostFloatMatrix @@ -94,7 +94,7 @@ float[][]\} on the heap. a host-memory CuVSMatrix -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:87`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:88`_ ### createHostByteMatrix @@ -105,7 +105,7 @@ static CuVSHostMatrix createHostByteMatrix(List data, int bytesPerVector Builds a host-memory CuVSMatrix from byte vectors without first materializing the list as an intermediate `byte[][]`. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:101`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:102`_ ### createHostByteMatrixFromArray @@ -115,7 +115,7 @@ static CuVSHostMatrix createHostByteMatrixFromArray(byte[][] data, int bytesPerV Builds a host-memory CuVSMatrix from a 2D byte array. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:112`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:113`_ ### nanosToMillis @@ -135,7 +135,7 @@ A utility method to convert nanoseconds to milliseconds. milliseconds -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:128`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:194`_ ### cuVSResourcesOrNull @@ -149,7 +149,7 @@ Creates an instance of CuVSResources. an instance of CuVSResources -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:137`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:203`_ ### handleThrowableWithIgnore @@ -172,7 +172,7 @@ A utility method that conditionally ignores certain throwable objects | --- | --- | | `IOException` | | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:165`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:231`_ ### info @@ -190,6 +190,6 @@ Utility to print info/debug messages via InfoStream. | `component` | the name of the index writer | | `msg` | the log message to push via the InfoStream | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:179`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:245`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:22`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:23`_ diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index 5320afe4af..27be666aac 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -132,7 +132,8 @@ public StandardDataset() {} } /** - * Invokes the native destroy_cagra_index to de-allocate the CAGRA index + * Invokes the native destroy_cagra_index to de-allocate the CAGRA index. Also attempts to close + * any dataset whose ownership transferred to this index during construction. */ @Override void close() throws Exception; @@ -428,6 +429,11 @@ interface Builder { /** * Sets the dataset for building the {@link CagraIndex}. * + *

The caller retains ownership until a build that uses this dataset returns successfully. + * The returned index then owns the dataset, and the caller must leave it open until the index + * is closed. If the build fails or uses another configured input source, ownership remains + * with the caller. + * * @param dataset a {@link CuVSMatrix} object containing the vectors * @return an instance of this Builder */ diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 52a28a0a6d..18461bbcfb 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -58,19 +58,12 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens // Create adjacency list for single node with no neighbors int[][] singleNodeAdjacency = new int[][] {{-1}}; // -1 indicates no neighbors - // Create CuVSMatrix from the adjacency list - CuVSMatrix adjacencyMatrix = CuVSMatrix.ofArray(singleNodeAdjacency); - - // Create layer data for single-level graph - List layerNodes = new ArrayList<>(); - List layerAdjacencies = new ArrayList<>(); - - // Layer 0: contains all nodes (just the single node) - layerNodes.add(null); // Layer 0 contains all nodes, so we don't need to store node list - layerAdjacencies.add(adjacencyMatrix); - - // Create the single-layer graph - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + // GPUBuiltHnswGraph copies the adjacency into heap-backed NeighborArrays. + try (CuVSMatrix adjacencyMatrix = CuVSMatrix.ofArray(singleNodeAdjacency)) { + List layerNodes = new ArrayList<>(); + layerNodes.add(null); // Layer 0 contains all nodes, so its node list is implicit. + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, List.of(adjacencyMatrix)); + } } /** @@ -88,6 +81,19 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( CagraIndexParams params, QuantizationType quantization) throws Throwable { + return createMultiLayerHnswGraph( + size, dimensions, adjacencyListMatrix, vectors, hnswLayers, params, quantization); + } + + private static GPUBuiltHnswGraph createMultiLayerHnswGraph( + int size, + int dimensions, + CuVSMatrix adjacencyListMatrix, + List vectors, + int hnswLayers, + CagraIndexParams params, + QuantizationType quantization) + throws Throwable { int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); @@ -102,70 +108,94 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( int currentLayerSize = size; int layerIndex = 1; Random random = new Random(); + Throwable failure = null; - while (layerIndex < hnswLayers && currentLayerSize > 1) { - // Calculate size for next layer (1/M of current layer) - int nextLayerSize = Math.max(2, currentLayerSize / M); - // Select nodes for this layer - SortedSet selectedNodesSet = new TreeSet<>(); - - if (layerIndex == 1) { - // Select from all nodes (Layer 0) - while (selectedNodesSet.size() < nextLayerSize) { - selectedNodesSet.add(random.nextInt(size)); - } - } else { - // Select from previous layer nodes - int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); - while (selectedNodesSet.size() < nextLayerSize) { - int idx = random.nextInt(prevLayerNodes.length); - selectedNodesSet.add(prevLayerNodes[idx]); + try { + while (layerIndex < hnswLayers && currentLayerSize > 1) { + // Calculate size for next layer (1/M of current layer) + int nextLayerSize = Math.max(2, currentLayerSize / M); + // Select nodes for this layer + SortedSet selectedNodesSet = new TreeSet<>(); + + if (layerIndex == 1) { + // Select from all nodes (Layer 0) + while (selectedNodesSet.size() < nextLayerSize) { + selectedNodesSet.add(random.nextInt(size)); + } + } else { + // Select from previous layer nodes + int[] prevLayerNodes = layerNodes.get(layerNodes.size() - 1); + while (selectedNodesSet.size() < nextLayerSize) { + int idx = random.nextInt(prevLayerNodes.length); + selectedNodesSet.add(prevLayerNodes[idx]); + } } - } - // Convert to sorted array - int[] selectedNodes = - selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray(); + // Convert to sorted array + int[] selectedNodes = + selectedNodesSet.stream().mapToInt(Integer::intValue).sorted().toArray(); - layerNodes.add(selectedNodes); + CuVSMatrix upperAdjacency; + if (quantization == QuantizationType.NONE) { + // Extract vectors for selected nodes + float[][] selectedVectors = new float[nextLayerSize][]; + for (int i = 0; i < nextLayerSize; i++) { + selectedVectors[i] = (float[]) vectors.get(selectedNodes[i]); + } - if (quantization == QuantizationType.NONE) { - // Extract vectors for selected nodes - float[][] selectedVectors = new float[nextLayerSize][]; - for (int i = 0; i < nextLayerSize; i++) { - selectedVectors[i] = (float[]) vectors.get(selectedNodes[i]); - } + // Build CAGRA graph for this layer + upperAdjacency = + buildCagraGraphForSubset( + selectedVectors, selectedNodes, 0, params, dimensions, quantization); - // Build CAGRA graph for this layer - layerAdjacencies.add( - buildCagraGraphForSubset( - selectedVectors, selectedNodes, 0, params, dimensions, quantization)); + } else { - } else { + // Extract vectors for selected nodes + int bytesPerVector = (dimensions + 7) / 8; + byte[][] selectedVectors = new byte[nextLayerSize][]; + for (int i = 0; i < nextLayerSize; i++) { + selectedVectors[i] = (byte[]) vectors.get(selectedNodes[i]); + } - // Extract vectors for selected nodes - int bytesPerVector = (dimensions + 7) / 8; - byte[][] selectedVectors = new byte[nextLayerSize][]; - for (int i = 0; i < nextLayerSize; i++) { - selectedVectors[i] = (byte[]) vectors.get(selectedNodes[i]); + // Build CAGRA graph for this layer + upperAdjacency = + buildCagraGraphForSubset( + selectedVectors, selectedNodes, bytesPerVector, params, dimensions, quantization); } - // Build CAGRA graph for this layer - layerAdjacencies.add( - buildCagraGraphForSubset( - selectedVectors, selectedNodes, bytesPerVector, params, dimensions, quantization)); - } + try { + // Register ownership before any later operation can fail. + layerAdjacencies.add(upperAdjacency); + } catch (Throwable registrationFailure) { + closeAfterFailure(upperAdjacency, registrationFailure); + throw registrationFailure; + } + layerNodes.add(selectedNodes); - // Update for next iteration - currentLayerSize = nextLayerSize; - layerIndex++; + // Update for next iteration + currentLayerSize = nextLayerSize; + layerIndex++; - // Use different seed for each layer - random = new Random(new Random().nextLong()); - } + // Use different seed for each layer + random = new Random(new Random().nextLong()); + } - // Create the multi-layer graph with all layers - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + // The graph eagerly copies all adjacency rows, so generated upper matrices can now close. + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies); + } catch (Throwable t) { + failure = t; + throw t; + } finally { + Throwable closeFailure = closeUpperLayerAdjacencies(layerAdjacencies); + if (closeFailure != null) { + if (failure == null) { + throw closeFailure; + } + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } + } + } } /** @@ -173,7 +203,6 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( * the Java heap. The list view copies only rows selected for an upper layer. */ static GPUBuiltHnswGraph createMultiLayerHnswGraph( - FieldInfo fieldInfo, int dimensions, CuVSMatrix adjacencyListMatrix, CuVSMatrix vectorDataset, @@ -206,14 +235,34 @@ public int size() { } }; return createMultiLayerHnswGraph( - fieldInfo, - size, - dimensions, - adjacencyListMatrix, - vectors, - hnswLayers, - params, - quantization); + size, dimensions, adjacencyListMatrix, vectors, hnswLayers, params, quantization); + } + + private static Throwable closeUpperLayerAdjacencies(List layerAdjacencies) { + Throwable failure = null; + // Layer 0 is borrowed from the outer CAGRA index. Only upper layers are owned here. + for (int i = layerAdjacencies.size() - 1; i >= 1; i--) { + try { + layerAdjacencies.get(i).close(); + } catch (Throwable closeFailure) { + if (failure == null) { + failure = closeFailure; + } else if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } + } + } + return failure; + } + + private static void closeAfterFailure(AutoCloseable resource, Throwable failure) { + try { + resource.close(); + } catch (Throwable closeFailure) { + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } + } } /** Builds a CAGRA graph for a selected vector subset. */ @@ -241,38 +290,36 @@ private static CuVSMatrix buildCagraGraphForSubset( private static CuVSMatrix buildCagraGraphForSubset( CuVSMatrix subsetDataset, int[] selectedNodes, CagraIndexParams params) throws Throwable { - - // Build CAGRA index for the subset - CagraIndex subsetIndex = - CagraIndex.newBuilder(getCuVSResourcesInstance()) - .withDataset(subsetDataset) - .withIndexParams(params) - .build(); - - // Get adjacency list from subset CAGRA index - CuVSMatrix cagraGraph = subsetIndex.getGraph(); - - long numNodes = cagraGraph.size(); - long degree = cagraGraph.columns(); - - // Create a re-mapped adjacency list - int[][] remappedAdjacency = new int[(int) numNodes][(int) degree]; - - for (int i = 0; i < numNodes; i++) { - RowView rv = cagraGraph.getRow(i); - for (int j = 0; j < degree && j < rv.size(); j++) { - int subsetIndex1 = rv.getAsInt(j); - // Map subset index to original node ID - if (subsetIndex1 >= 0 && subsetIndex1 < selectedNodes.length) { - remappedAdjacency[i][j] = selectedNodes[subsetIndex1]; - } else { - // Invalid index, use self-reference - remappedAdjacency[i][j] = selectedNodes[i]; + int[][] remappedAdjacency; + try (Utils.OwnedIndex ownedIndex = Utils.ownDataset(subsetDataset)) { + CagraIndex subsetIndex = + CagraIndex.newBuilder(getCuVSResourcesInstance()) + .withDataset(subsetDataset) + .withIndexParams(params) + .build(); + ownedIndex.transferTo(subsetIndex); + + CuVSMatrix cagraGraph = subsetIndex.getGraph(); + long numNodes = cagraGraph.size(); + long degree = cagraGraph.columns(); + remappedAdjacency = new int[(int) numNodes][(int) degree]; + + for (int i = 0; i < numNodes; i++) { + RowView rv = cagraGraph.getRow(i); + for (int j = 0; j < degree && j < rv.size(); j++) { + int subsetIndex1 = rv.getAsInt(j); + // Map subset index to original node ID + if (subsetIndex1 >= 0 && subsetIndex1 < selectedNodes.length) { + remappedAdjacency[i][j] = selectedNodes[subsetIndex1]; + } else { + // Invalid index, use self-reference + remappedAdjacency[i][j] = selectedNodes[i]; + } } } } - subsetIndex.close(); + // Build the returned matrix only after the subset index and its dataset have closed. return CuVSMatrix.ofArray(remappedAdjacency); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index 7e9f888e32..b7a50316b4 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -75,7 +75,10 @@ private NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size) { if (rv != null && rv.size() > 0) { neighbors[i] = new NeighborArray((int) rv.size(), true); for (int j = 0; j < rv.size(); j++) { - neighbors[i].addInOrder(rv.getAsInt(j), 1.0f - (j * 0.001f)); + int neighbor = rv.getAsInt(j); + if (neighbor >= 0) { + neighbors[i].addInOrder(neighbor, 1.0f - (j * 0.001f)); + } } } else { neighbors[i] = new NeighborArray(0, true); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index bbfd4309b3..bbafbe827a 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -166,7 +166,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro * @throws IOException */ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { - try { + try (Utils.OwnedIndex ownedIndex = Utils.ownDataset(dataset)) { int size = (int) dataset.size(); CagraIndexParams params = CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); @@ -175,11 +175,11 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw .withDataset(dataset) .withIndexParams(params) .build(); + ownedIndex.transferTo(cagraIndex); CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); int dimensions = fieldInfo.getVectorDimension(); GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( - fieldInfo, dimensions, adjacencyListMatrix, dataset, @@ -198,9 +198,8 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw size, hnswGraph, graphLevelNodeOffsets); - cagraIndex.close(); } catch (Throwable t) { - Utils.handleThrowable(t); + throw Utils.handleThrowable(t); } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 61cc31545a..9efcbec3e2 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -146,38 +146,38 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException * @throws IOException */ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { - if (vectors.size() == 0) { - writeEmpty(fieldInfo, hnswMeta); + int size = vectors.size(); + if (writeTrivialField(fieldInfo, size)) { return; } try { int dimensions = fieldInfo.getVectorDimension(); int bytesPerVector = (dimensions + 7) / 8; - CuVSMatrix dataset = Utils.createHostByteMatrix(vectors, bytesPerVector); + writeNonTrivialField(fieldInfo, dataset); + } catch (Throwable t) { + throw Utils.handleThrowable(t); + } + } - if (dataset.size() < 2) { - writeSingleVectorGraph(fieldInfo, vectors); - return; - } - + /** Builds and writes an index from an owned binary-vector matrix. */ + private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { + try (Utils.OwnedIndex ownedIndex = Utils.ownDataset(dataset)) { + int size = (int) dataset.size(); + int dimensions = fieldInfo.getVectorDimension(); CagraIndexParams params = CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); - CagraIndex cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()) .withDataset(dataset) .withIndexParams(params) .build(); + ownedIndex.transferTo(cagraIndex); CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); - int size = (int) dataset.size(); - - // Create multi-layer HNSW graph from CAGRA GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( - fieldInfo, dimensions, adjacencyListMatrix, dataset, @@ -186,11 +186,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw QuantizationType.BINARY); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - // Write the graph to the vector index int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - - // Write metadata writeMeta( hnswVectorIndex, hnswMeta, @@ -200,14 +197,24 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw size, hnswGraph, graphLevelNodeOffsets); - - cagraIndex.close(); - } catch (Throwable t) { - Utils.handleThrowable(t); + throw Utils.handleThrowable(t); } } + /** Writes the empty or one-vector representation, if {@code size} is trivial. */ + private boolean writeTrivialField(FieldInfo fieldInfo, int size) throws IOException { + if (size == 0) { + writeEmpty(fieldInfo, hnswMeta); + return true; + } + if (size == 1) { + writeSingleVectorGraph(fieldInfo); + return true; + } + return false; + } + /** * Build the indexes and writes it to the disk. */ @@ -259,11 +266,9 @@ private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) thr * Builds and writes a single vector graph. * * @param fieldInfo instance of FieldInfo - * @param vectors the list of binary quantized vectors * @throws IOException I/O Exceptions */ - private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) - throws IOException { + private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { // Workaround for CAGRA not supporting single vector indexes try { int size = 1; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index ab76df4766..60cc338e4d 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -166,8 +166,8 @@ private static byte[] convertSignedToUnsigned(byte[] signedVector) { * @throws IOException */ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { - if (vectors.size() == 0) { - writeEmpty(fieldInfo, hnswMeta); + int size = vectors.size(); + if (writeTrivialField(fieldInfo, size)) { return; } @@ -182,27 +182,29 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE // Create CuVSMatrix with BYTE data type (unsigned bytes) CuVSMatrix dataset = Utils.createHostByteMatrix(unsignedVectors, dimensions); + writeNonTrivialField(fieldInfo, dataset); + } catch (Throwable t) { + throw Utils.handleThrowable(t); + } + } - if (dataset.size() < 2) { - writeSingleVectorGraph(fieldInfo, unsignedVectors); - return; - } - + /** Builds and writes an index from an owned scalar-vector matrix. */ + private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throws IOException { + try (Utils.OwnedIndex ownedIndex = Utils.ownDataset(dataset)) { + int size = (int) dataset.size(); + int dimensions = fieldInfo.getVectorDimension(); CagraIndexParams params = CagraIndexParamsFactory.create(acceleratedHNSWParams, dataset.size(), dataset.columns()); - CagraIndex cagraIndex = CagraIndex.newBuilder(getCuVSResourcesInstance()) .withDataset(dataset) .withIndexParams(params) .build(); + ownedIndex.transferTo(cagraIndex); CuVSMatrix adjacencyListMatrix = cagraIndex.getGraph(); - - int size = (int) dataset.size(); GPUBuiltHnswGraph hnswGraph = createMultiLayerHnswGraph( - fieldInfo, dimensions, adjacencyListMatrix, dataset, @@ -211,13 +213,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE QuantizationType.SCALAR); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - - // Write the graph to the vector index int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); - long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; - - // Write metadata writeMeta( hnswVectorIndex, hnswMeta, @@ -227,13 +224,24 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE size, hnswGraph, graphLevelNodeOffsets); - - cagraIndex.close(); } catch (Throwable t) { - Utils.handleThrowable(t); + throw Utils.handleThrowable(t); } } + /** Writes the empty or one-vector representation, if {@code size} is trivial. */ + private boolean writeTrivialField(FieldInfo fieldInfo, int size) throws IOException { + if (size == 0) { + writeEmpty(fieldInfo, hnswMeta); + return true; + } + if (size == 1) { + writeSingleVectorGraph(fieldInfo); + return true; + } + return false; + } + /** * Build the indexes and writes it to the disk. */ @@ -285,11 +293,9 @@ private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) thr * Builds and writes a single vector graph. * * @param fieldInfo instance of FieldInfo - * @param vectors the list of scalar quantized vectors (already converted to unsigned) * @throws IOException I/O Exceptions */ - private void writeSingleVectorGraph(FieldInfo fieldInfo, List vectors) - throws IOException { + private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { // Workaround for CAGRA not supporting single vector indexes try { int size = 1; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java index 5547db2bb4..4d348f6d75 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -10,6 +10,7 @@ import java.io.IOException; import java.time.Duration; import java.util.List; +import java.util.Objects; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.util.InfoStream; @@ -120,6 +121,71 @@ static CuVSHostMatrix createHostByteMatrixFromArray(byte[][] data, int bytesPerV } } + // Starts an ownership scope for a dataset that may later be transferred to an index. + static OwnedIndex ownDataset(AutoCloseable dataset) { + return new OwnedIndex<>(dataset); + } + + /** + * Owns a dataset until {@link #transferTo} records a successful index build. Closing the scope + * then closes either the original dataset or the index that owns it. + */ + static final class OwnedIndex implements AutoCloseable { + private AutoCloseable dataset; + private I index; + private boolean closed; + + private OwnedIndex(AutoCloseable dataset) { + this.dataset = Objects.requireNonNull(dataset, "dataset"); + } + + void transferTo(I index) { + if (closed || this.index != null) { + throw new IllegalStateException("Dataset ownership has already been transferred"); + } + this.index = Objects.requireNonNull(index, "index"); + } + + @Override + public void close() throws Exception { + if (closed) { + return; + } + closed = true; + AutoCloseable ownedDataset = dataset; + I ownedIndex = index; + dataset = null; + index = null; + if (ownedIndex == null) { + ownedDataset.close(); + return; + } + + try { + ownedIndex.close(); + } catch (Throwable indexCloseFailure) { + try { + ownedDataset.close(); + } catch (Throwable datasetCloseFailure) { + if (indexCloseFailure != datasetCloseFailure) { + indexCloseFailure.addSuppressed(datasetCloseFailure); + } + } + rethrowCloseFailure(indexCloseFailure); + } + } + } + + private static void rethrowCloseFailure(Throwable failure) throws Exception { + if (failure instanceof Exception exception) { + throw exception; + } + if (failure instanceof Error error) { + throw error; + } + throw new AssertionError("Unexpected throwable from AutoCloseable.close()", failure); + } + /** * A utility method to convert nanoseconds to milliseconds. * diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java index 28fa4d974e..f748f14f69 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWDeletedDocuments.java @@ -455,6 +455,11 @@ private void assertTrivialLiveVectorMerge(int liveVectors) throws Exception { assertEquals(liveVectors, graph.size()); assertEquals(liveVectors == 0 ? 0 : 1, graph.numLevels()); assertEquals(liveVectors, graph.getNodesOnLevel(0).size()); + if (liveVectors == 1) { + assertEquals(0, graph.maxConn()); + graph.seek(0, 0); + assertEquals(NO_MORE_DOCS, graph.nextNeighbor()); + } ((CodecReader) leaf).getVectorReader().checkIntegrity(); IndexSearcher searcher = new IndexSearcher(reader); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUpperLayers.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUpperLayers.java index 0c7b5db007..75d4e43602 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUpperLayers.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWUpperLayers.java @@ -4,6 +4,8 @@ */ package com.nvidia.cuvs.lucene; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; @@ -35,7 +37,6 @@ public void testLegacyListOverloadDescriptorIsPresent() throws Exception { Method matrixOverload = AcceleratedHNSWUtils.class.getDeclaredMethod( "createMultiLayerHnswGraph", - FieldInfo.class, int.class, CuVSMatrix.class, CuVSMatrix.class, @@ -44,4 +45,15 @@ public void testLegacyListOverloadDescriptorIsPresent() throws Exception { QuantizationType.class); assertFalse(Modifier.isPublic(matrixOverload.getModifiers())); } + + @Test + public void testSingleVectorGraphHasNoNeighbors() throws Throwable { + GPUBuiltHnswGraph graph = AcceleratedHNSWUtils.createSingleVectorHnswGraph(1, 32); + + assertEquals(1, graph.numLevels()); + assertEquals(0, graph.maxConn()); + assertEquals(0, graph.getNeighbors(0, 0).size()); + graph.seek(0, 0); + assertEquals(NO_MORE_DOCS, graph.nextNeighbor()); + } } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java index 6aead1bd8d..c010b04e12 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtilsThrowableHandling.java @@ -47,4 +47,70 @@ public void testHandleThrowableWrapsCheckedExceptionWithCause() { assertSame(exception, thrown.getCause()); } + + @Test + public void testOwnedIndexClosesUntransferredDatasetOnce() throws Exception { + TrackingCloseable dataset = new TrackingCloseable(null); + Utils.OwnedIndex owned = Utils.ownDataset(dataset); + + owned.close(); + owned.close(); + + assertEquals(1, dataset.closeCount); + } + + @Test + public void testOwnedIndexClosesTransferredIndexInsteadOfDataset() throws Exception { + TrackingCloseable dataset = new TrackingCloseable(null); + TrackingCloseable index = new TrackingCloseable(null); + Utils.OwnedIndex owned = Utils.ownDataset(dataset); + owned.transferTo(index); + + owned.close(); + + assertEquals(1, index.closeCount); + assertEquals(0, dataset.closeCount); + } + + @Test + public void testOwnedIndexPreservesBodyAndCleanupFailures() { + IOException bodyFailure = new IOException("body"); + IOException indexCloseFailure = new IOException("index close"); + IOException datasetCloseFailure = new IOException("dataset close"); + TrackingCloseable dataset = new TrackingCloseable(datasetCloseFailure); + TrackingCloseable index = new TrackingCloseable(indexCloseFailure); + + IOException thrown = + assertThrows( + IOException.class, + () -> { + try (Utils.OwnedIndex owned = Utils.ownDataset(dataset)) { + owned.transferTo(index); + throw bodyFailure; + } + }); + + assertSame(bodyFailure, thrown); + assertArrayEquals(new Throwable[] {indexCloseFailure}, thrown.getSuppressed()); + assertArrayEquals(new Throwable[] {datasetCloseFailure}, indexCloseFailure.getSuppressed()); + assertEquals(1, index.closeCount); + assertEquals(1, dataset.closeCount); + } + + private static final class TrackingCloseable implements AutoCloseable { + private final Exception failure; + private int closeCount; + + private TrackingCloseable(Exception failure) { + this.failure = failure; + } + + @Override + public void close() throws Exception { + closeCount++; + if (failure != null) { + throw failure; + } + } + } } From f72e26de880285dbacac7329319e0139658e8ea2 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Tue, 29 Sep 2026 21:04:19 +0000 Subject: [PATCH 14/21] Harden parallel HNSW graph processing Add independently configurable graph workers, bounded shared execution, physical-memory-aware copy admission, and byte-bounded serialization waves. Cover persistence, lifecycle, failure, concurrency, and high-degree serialization behavior. --- ...vidia-cuvs-lucene-acceleratedhnswparams.md | 103 ++-- ...nvidia-cuvs-lucene-acceleratedhnswutils.md | 20 +- ...om-nvidia-cuvs-lucene-gpubuilthnswgraph.md | 44 +- ...leratedhnswbinaryquantizedvectorswriter.md | 10 +- ...leratedhnswscalarquantizedvectorswriter.md | 10 +- fern/pages/user_guide/lucene.md | 4 +- .../cuvs/lucene/AcceleratedHNSWParams.java | 48 +- .../cuvs/lucene/AcceleratedHNSWUtils.java | 170 ++++--- .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 166 +++---- .../cuvs/lucene/GraphCopyMemoryBudget.java | 137 ++++++ .../nvidia/cuvs/lucene/GraphWorkExecutor.java | 167 +++++++ .../Lucene99AcceleratedHNSWVectorsWriter.java | 6 +- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 7 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 7 +- .../cuvs/lucene/IntGraphTestMatrix.java | 238 +++++++++ .../lucene/TestAcceleratedHNSWParams.java | 26 + .../lucene/TestCagraIndexParamsFactory.java | 1 + .../lucene/TestGraphCopyMemoryBudget.java | 152 ++++++ .../TestGraphThreadsPersistedIndex.java | 127 +++++ .../cuvs/lucene/TestGraphWorkExecutor.java | 455 ++++++++++++++++++ .../TestParallelGraphMaterialization.java | 153 ++++++ .../TestParallelGraphSerialization.java | 192 ++++++++ 22 files changed, 1950 insertions(+), 293 deletions(-) create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphWorkExecutor.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/IntGraphTestMatrix.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphThreadsPersistedIndex.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphWorkExecutor.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md index d3147c842e..f4ca78f79f 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md @@ -18,14 +18,28 @@ public class AcceleratedHNSWParams public int getWriterThreads() ``` -Get the maximum thread count for cuVS writes and accelerated-HNSW graph materialization and -serialization. +Get the native cuVS writer threads parameter. **Returns** -maximum writer and graph-processing thread count +cuVS writer threads parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:148`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:154`_ + +### getGraphThreads + +```java +public int getGraphThreads() +``` + +Get the maximum threads per HNSW graph materialization or serialization operation. The count +includes the calling thread; shared helper capacity may reduce actual concurrency. + +**Returns** + +HNSW graph processing threads parameter + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:164`_ ### getIntermediateGraphDegree @@ -39,7 +53,7 @@ Get the intermediate graph degree the graph degree parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:157`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:173`_ ### getGraphdegree @@ -53,7 +67,7 @@ Get the graph degree the graph degree parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:166`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:182`_ ### getHnswLayers @@ -67,7 +81,7 @@ Get the number of HNSW layers the number of HNSW layers -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:175`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:191`_ ### getMaxConn @@ -81,7 +95,7 @@ Get the max connection parameter the max connection parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:184`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:200`_ ### getBeamWidth @@ -95,7 +109,7 @@ Get the beam width parameter the beam width parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:193`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:209`_ ### getCagraGraphBuildAlgo @@ -109,7 +123,7 @@ Get the CAGRA graph build algorithm the CAGRA graph build algorithm -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:202`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:218`_ ### getCuVSIvfPqParams @@ -123,7 +137,7 @@ Get the instance of `CuVSIvfPqParams` the instance of `CuVSIvfPqParams` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:211`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:227`_ ### getNumMergeWorkers @@ -137,7 +151,7 @@ Get the number of merge workers set to be used in the fallback mechanism the number of merge workers -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:220`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:236`_ ### getMergeExec @@ -151,7 +165,7 @@ Get the instance of the `ExecutorService` to be used in the fallback mechanism the instance of the `ExecutorService` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:229`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:245`_ ### getStrategy @@ -168,7 +182,7 @@ When CUSTOM is chosen, the build algorithm and its parameters (either defaults o get the chosen `Strategy` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:241`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:257`_ ### getCuvsDistanceType @@ -182,7 +196,7 @@ Get the cuvs distance type the distance type -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:250`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:266`_ ### getNNDescentNumIterations @@ -196,7 +210,7 @@ get the number of Iterations to run if building with NN_DESCENT the number of iterations for NN_DESCENT -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:259`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:275`_ ### getHnswHeuristicType @@ -211,7 +225,7 @@ beamWidth. Only consulted under the `Strategy#HEURISTIC` strategy. the `HnswHeuristicType` to hand to cuVS -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:269`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:285`_ ### withWriterThreads @@ -219,8 +233,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWP public Builder withWriterThreads(int writerThreads) ``` -Set the maximum number of threads used for cuVS writes and accelerated-HNSW graph -materialization and serialization. +Set the number of native cuVS writer threads while building the index. Valid range - Minimum: \{@value MIN_WRITER_THREADS\}, Maximum: \{@value MAX_WRITER_THREADS\} Default value - \{@value DEFAULT_WRITER_THREADS\} @@ -234,7 +247,29 @@ Default value - \{@value DEFAULT_WRITER_THREADS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:335`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:353`_ + +### withGraphThreads + +```java +public Builder withGraphThreads(int graphThreads) +``` + +Set the maximum threads per HNSW graph materialization or serialization operation. The count +includes the calling thread. Valid range - Minimum: \{@value MIN_GRAPH_THREADS\}, Maximum: +\{@value MAX_GRAPH_THREADS\}. Default value - \{@value DEFAULT_GRAPH_THREADS\}. + +**Parameters** + +| Name | Description | +| --- | --- | +| `graphThreads` | maximum graph-processing threads per operation | + +**Returns** + +instance of `Builder` + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:366`_ ### withIntermediateGraphDegree @@ -256,7 +291,7 @@ Default value - \{@value DEFAULT_INT_GRAPH_DEGREE\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:348`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:379`_ ### withGraphDegree @@ -278,7 +313,7 @@ Default value - \{@value DEFAULT_GRAPH_DEGREE\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:361`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:392`_ ### withHNSWLayer @@ -300,7 +335,7 @@ Default value - \{@value DEFAULT_HNSW_LAYERS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:374`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:405`_ ### withMaxConn @@ -322,7 +357,7 @@ Default value - \{@value DEFAULT_MAX_CONN\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:387`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:418`_ ### withBeamWidth @@ -344,7 +379,7 @@ Default value - \{@value DEFAULT_BEAM_WIDTH\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:400`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:431`_ ### withCagraGraphBuildAlgo @@ -365,7 +400,7 @@ Default value - NN_DESCENT instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:412`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:443`_ ### withCuVSIvfPqParams @@ -385,7 +420,7 @@ Set the instance of `CuVSIvfPqParams` instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:423`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:454`_ ### withNumMergeWorkers @@ -406,7 +441,7 @@ Default value - \{@value DEFAULT_NUM_MERGE_WORKERS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:435`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:466`_ ### withMergeExecutorService @@ -427,7 +462,7 @@ Default value an instance with one thread instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:447`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:478`_ ### withStrategy @@ -453,7 +488,7 @@ Default value - HEURISTIC instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:464`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:495`_ ### withCuvsDistanceType @@ -473,7 +508,7 @@ Set the CuvsDistanceType instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:475`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:506`_ ### withNNDescentNumIterations @@ -496,7 +531,7 @@ Default value - \{@value DEFAULT_NN_DESCENT_NUM_ITERATIONS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:489`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:520`_ ### withHnswHeuristicType @@ -520,7 +555,7 @@ the equivalent HNSW graph (graph degree = 2 * maxConn). instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:504`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:535`_ ### build @@ -534,6 +569,6 @@ Create an instance of `AcceleratedHNSWParams` instance of `AcceleratedHNSWParams` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:549`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:582`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:17`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md index 640f8bdaeb..f9321443e4 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md @@ -21,7 +21,7 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens Creates a dummy HNSW graph for a single vector. The graph will have 1 level with 1 node and no neighbors. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:61`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:58`_ ### createMultiLayerHnswGraph @@ -33,7 +33,7 @@ Creates up to `hnswLayers` total layers. Layer 0 uses the full CAGRA graph. Each layer samples `max(2, floor(previousLayerSize / M))` nodes. The value `M` is the ceiling of half the layer-0 graph degree. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:79`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:76`_ ### createMultiLayerHnswGraph @@ -44,7 +44,7 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( int dimensions, CuVSMatrix a Creates a multi-layer HNSW graph from a native matrix without copying the complete dataset to the Java heap. The list view copies only rows selected for an upper layer. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:211`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:215`_ ### writeGraph @@ -71,7 +71,7 @@ a 2D array of offsets | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:360`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:370`_ ### writeMeta @@ -100,7 +100,7 @@ Writes the meta information for the index. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:524`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:516`_ ### printInfoStream @@ -116,7 +116,7 @@ A utility method to print info/debugging messages using InfoStream. | --- | --- | | `msg` | the debugging message to print | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:606`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:598`_ ### writeEmpty @@ -138,7 +138,7 @@ Writes an empty meta information for the field. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:618`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:610`_ ### quantizeFloatVectorsToBinary @@ -161,7 +161,7 @@ Bits are packed: 8 dimensions per byte. A list of byte binary representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:631`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:623`_ ### quantizeFloatVectorsToScalar @@ -181,6 +181,6 @@ Scalar quantization. A list of byte scalar representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:673`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:665`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:38`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:35`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md index 3367485397..239e1a3f33 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md @@ -31,27 +31,7 @@ Multi-layer constructor that supports arbitrary number of layers. | `layerNodes` | the nodes on the layer | | `layerAdjacencies` | adjacency list | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:55`_ - -### GPUBuiltHnswGraph - -```java -public GPUBuiltHnswGraph( int size, int dimensions, List layerNodes, List layerAdjacencies, int numThreads) throws IOException -``` - -Builds a graph while materializing adjacency rows with the requested number of threads. - -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:61`_ - -### fitsParallelGraphCopyBudget - -```java -static boolean fitsParallelGraphCopyBudget(long rows, long columns) -``` - -Returns whether an INT32 adjacency can be copied without exceeding the native-host budget. - -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:169`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:56`_ ### getNodesOnLevel @@ -61,7 +41,7 @@ public NodesIterator getNodesOnLevel(int level) Get all nodes on a given level as node 0th ordinals. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:256`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:228`_ ### getNeighbors @@ -82,7 +62,7 @@ Get the neighbors for the node and the level it resides. an instance of NeighborArray -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:274`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:246`_ ### seek @@ -92,7 +72,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Move the pointer to exactly the given level's target. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:299`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:271`_ ### nextNeighbor @@ -102,7 +82,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Iterates over the neighbor list. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:309`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:281`_ ### entryNode @@ -112,7 +92,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns graph's entry point on the top level. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:340`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:312`_ ### maxConn @@ -122,7 +102,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap returns M, the maximum number of connections for a node. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:359`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:331`_ ### neighborCount @@ -132,7 +112,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns the neighbor count. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:374`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:346`_ ### size @@ -142,7 +122,7 @@ public int size() Returns the number of nodes in the graph. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:449`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:421`_ ### numLevels @@ -156,7 +136,7 @@ Returns the number of levels in the HNSW graph. the number of levels -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:458`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:430`_ ### dimensions @@ -170,6 +150,6 @@ Gets the vector dimension. the vector dimension -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:467`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:439`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:29`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:27`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md index ae41d679dd..b355a3784f 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:223`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:222`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:309`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:308`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:340`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:339`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:360`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:359`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:369`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:368`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:57`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md index 8a67d9e920..6ba5228a80 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:250`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:249`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:334`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:333`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:365`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:364`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:385`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:384`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:394`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:393`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:56`_ diff --git a/fern/pages/user_guide/lucene.md b/fern/pages/user_guide/lucene.md index 2f3e49cd20..670a67305f 100644 --- a/fern/pages/user_guide/lucene.md +++ b/fern/pages/user_guide/lucene.md @@ -231,9 +231,7 @@ Both `AcceleratedHNSWParams` and `GPUSearchParams` default to a `HEURISTIC` stra - For the accelerated HNSW codecs, set `maxConn` and `beamWidth`, the HNSW parameters you would tune on the CPU. cuVS derives graph degrees and the build algorithm from them. These two values also configure the CPU fallback writer, so one setting covers both paths. - For the GPU search codec, set `buildQuality`. Higher values spend more build time for a higher-quality graph. This codec also passes `graphDegree` into the heuristic, so leave it at its default unless you intend to cap the graph. -`writerThreads` controls native cuVS writer concurrency for both parameter types. For the accelerated HNSW codecs, it also sets the maximum number of CPU threads used to materialize and serialize the finished graph; the default of one keeps this post-build processing serial. The GPU search codec continues to use the setting only for native writer concurrency and defaults to 32. - -With more than one writer thread, parallel adjacency materialization is eligible for layers with at least 65,536 nodes. A device adjacency is copied once to native host memory when that copy is at most 4 GiB; larger device adjacencies retain serial row access. Level-zero serialization is independently eligible at 65,536 nodes and runs in fixed waves of at most 1,048,576 nodes; upper levels remain serial. +`writerThreads` controls native cuVS build concurrency. The accelerated HNSW codecs default to one writer thread, while the GPU search codec defaults to 32. Accelerated HNSW also offers `graphThreads` for the later CPU graph-materialization and serialization stages; it defaults to one and counts the calling thread. Parallel materialization of a device-backed graph uses a temporary host copy only when the JVM reports enough free physical memory for that copy, the Lucene graph, and an adjacency-sized safety allowance. Otherwise, materialization remains serial. Switching either class to the `CUSTOM` strategy exposes the underlying CAGRA parameters directly, including `graphDegree`, `intermediateGraphDegree`, the graph build algorithm, and its parameters. Use `CUSTOM` only when you have measurements that justify specific values; the defaults derived by cuVS are a better starting point. For background on the parameters themselves, see the [CAGRA indexing guide](/user-guide/api-guides/indexing-guide/cagra) and the [tuning guide](/getting-started/introduction/tuning-indexes). diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index 458566931f..1f8c13609b 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -34,6 +34,8 @@ public static enum Strategy { // Bounds for the public CAGRA and HNSW build parameters. public static final int MIN_WRITER_THREADS = 1; public static final int MAX_WRITER_THREADS = 512; + public static final int MIN_GRAPH_THREADS = 1; + public static final int MAX_GRAPH_THREADS = 512; public static final int MIN_INT_GRAPH_DEG = 2; public static final int MAX_INT_GRAPH_DEG = 512; public static final int MIN_GRAPH_DEG = 1; @@ -50,6 +52,7 @@ public static enum Strategy { public static final int MAX_NN_DESCENT_NUM_ITERATIONS = 100; public static final int DEFAULT_WRITER_THREADS = 1; + public static final int DEFAULT_GRAPH_THREADS = 1; public static final int DEFAULT_INT_GRAPH_DEGREE = 128; public static final int DEFAULT_GRAPH_DEGREE = 64; public static final int DEFAULT_HNSW_LAYERS = 1; @@ -75,6 +78,7 @@ public static enum Strategy { }; private final int writerThreads; + private final int graphThreads; private final int intermediateGraphDegree; private final int graphdegree; private final int hnswLayers; @@ -92,8 +96,9 @@ public static enum Strategy { /** * Constructs an instance of {@link AcceleratedHNSWParams} with specific parameter values. * - * @param writerThreads Maximum number of threads to use for cuVS writes and accelerated-HNSW - * graph materialization and serialization. + * @param writerThreads Number of native cuVS writer threads to use. + * @param graphThreads Maximum threads per HNSW graph materialization or serialization operation, + * including the calling thread. * @param intermediateGraphDegree The intermediate graph degree while building the CAGRA index. * @param graphdegree The graph degree to use while building the CAGRA index. * @param hnswLayers The number of HNSW layers to build in the HNSW index. @@ -110,6 +115,7 @@ public static enum Strategy { */ private AcceleratedHNSWParams( int writerThreads, + int graphThreads, int intermediateGraphDegree, int graphdegree, int hnswLayers, @@ -125,6 +131,7 @@ private AcceleratedHNSWParams( HnswHeuristicType hnswHeuristicType) { super(); this.writerThreads = writerThreads; + this.graphThreads = graphThreads; this.intermediateGraphDegree = intermediateGraphDegree; this.graphdegree = graphdegree; this.hnswLayers = hnswLayers; @@ -141,15 +148,24 @@ private AcceleratedHNSWParams( } /** - * Get the maximum thread count for cuVS writes and accelerated-HNSW graph materialization and - * serialization. + * Get the native cuVS writer threads parameter. * - * @return maximum writer and graph-processing thread count + * @return cuVS writer threads parameter */ public int getWriterThreads() { return writerThreads; } + /** + * Get the maximum threads per HNSW graph materialization or serialization operation. The count + * includes the calling thread; shared helper capacity may reduce actual concurrency. + * + * @return HNSW graph processing threads parameter + */ + public int getGraphThreads() { + return graphThreads; + } + /** * Get the intermediate graph degree * @@ -275,6 +291,8 @@ public HnswHeuristicType getHnswHeuristicType() { public String toString() { return "AcceleratedHNSWParams [writerThreads=" + writerThreads + + ", graphThreads=" + + graphThreads + ", intermediateGraphDegree=" + intermediateGraphDegree + ", graphdegree=" @@ -310,6 +328,7 @@ public String toString() { public static class Builder { private int writerThreads = DEFAULT_WRITER_THREADS; + private int graphThreads = DEFAULT_GRAPH_THREADS; private int intermediateGraphDegree = DEFAULT_INT_GRAPH_DEGREE; private int graphdegree = DEFAULT_GRAPH_DEGREE; private int hnswLayers = DEFAULT_HNSW_LAYERS; @@ -325,8 +344,7 @@ public static class Builder { private HnswHeuristicType hnswHeuristicType = DEFAULT_HNSW_HEURISTIC_TYPE; /** - * Set the maximum number of threads used for cuVS writes and accelerated-HNSW graph - * materialization and serialization. + * Set the number of native cuVS writer threads while building the index. * Valid range - Minimum: {@value MIN_WRITER_THREADS}, Maximum: {@value MAX_WRITER_THREADS} * Default value - {@value DEFAULT_WRITER_THREADS} * @@ -338,6 +356,19 @@ public Builder withWriterThreads(int writerThreads) { return this; } + /** + * Set the maximum threads per HNSW graph materialization or serialization operation. The count + * includes the calling thread. Valid range - Minimum: {@value MIN_GRAPH_THREADS}, Maximum: + * {@value MAX_GRAPH_THREADS}. Default value - {@value DEFAULT_GRAPH_THREADS}. + * + * @param graphThreads maximum graph-processing threads per operation + * @return instance of {@link Builder} + */ + public Builder withGraphThreads(int graphThreads) { + this.graphThreads = graphThreads; + return this; + } + /** * Set the intermediate graph degree to use while building CAGRA index * Valid range - Minimum: {@value MIN_INT_GRAPH_DEG}, Maximum: {@value MAX_INT_GRAPH_DEG} @@ -515,6 +546,8 @@ public Builder withHnswHeuristicType(HnswHeuristicType hnswHeuristicType) { private void validate() throws IllegalArgumentException { ParameterValidation.checkRange( "writerThreads", writerThreads, MIN_WRITER_THREADS, MAX_WRITER_THREADS); + ParameterValidation.checkRange( + "graphThreads", graphThreads, MIN_GRAPH_THREADS, MAX_GRAPH_THREADS); ParameterValidation.checkRange( "intermediateGraphDegree", intermediateGraphDegree, MIN_INT_GRAPH_DEG, MAX_INT_GRAPH_DEG); ParameterValidation.checkRange("graphdegree", graphdegree, MIN_GRAPH_DEG, MAX_GRAPH_DEG); @@ -557,6 +590,7 @@ public AcceleratedHNSWParams build() { validate(); return new AcceleratedHNSWParams( writerThreads, + graphThreads, intermediateGraphDegree, graphdegree, hnswLayers, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index ae5d9b1714..620e4464f6 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -21,11 +21,8 @@ import java.util.SortedSet; import java.util.TreeSet; import java.util.concurrent.Callable; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.search.TaskExecutor; import org.apache.lucene.store.ByteBuffersDataOutput; import org.apache.lucene.store.DataOutput; import org.apache.lucene.store.IndexOutput; @@ -68,7 +65,7 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens try (CuVSMatrix adjacencyMatrix = CuVSMatrix.ofArray(singleNodeAdjacency)) { List layerNodes = new ArrayList<>(); layerNodes.add(null); // Layer 0 contains all nodes, so its node list is implicit. - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, List.of(adjacencyMatrix), 1); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, List.of(adjacencyMatrix)); } } @@ -88,7 +85,14 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( QuantizationType quantization) throws Throwable { return createMultiLayerHnswGraph( - size, dimensions, adjacencyListMatrix, vectors, hnswLayers, params, quantization, 1); + size, + dimensions, + adjacencyListMatrix, + vectors, + hnswLayers, + params, + quantization, + AcceleratedHNSWParams.DEFAULT_GRAPH_THREADS); } private static GPUBuiltHnswGraph createMultiLayerHnswGraph( @@ -99,7 +103,7 @@ private static GPUBuiltHnswGraph createMultiLayerHnswGraph( int hnswLayers, CagraIndexParams params, QuantizationType quantization, - int numThreads) + int graphThreads) throws Throwable { int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); @@ -188,7 +192,7 @@ private static GPUBuiltHnswGraph createMultiLayerHnswGraph( } // The graph eagerly copies all adjacency rows, so generated upper matrices can now close. - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, numThreads); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, graphThreads); } catch (Throwable t) { failure = t; throw t; @@ -218,7 +222,13 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( QuantizationType quantization) throws Throwable { return createMultiLayerHnswGraph( - dimensions, adjacencyListMatrix, vectorDataset, hnswLayers, params, quantization, 1); + dimensions, + adjacencyListMatrix, + vectorDataset, + hnswLayers, + params, + quantization, + AcceleratedHNSWParams.DEFAULT_GRAPH_THREADS); } static GPUBuiltHnswGraph createMultiLayerHnswGraph( @@ -228,7 +238,7 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( int hnswLayers, CagraIndexParams params, QuantizationType quantization, - int numThreads) + int graphThreads) throws Throwable { int size = Math.toIntExact(vectorDataset.size()); // Matrix columns are the stored width: binary vectors are bit-packed, while scalar and float @@ -262,7 +272,7 @@ public int size() { hnswLayers, params, quantization, - numThreads); + graphThreads); } private static Throwable closeUpperLayerAdjacencies(List layerAdjacencies) { @@ -360,27 +370,21 @@ private static CuVSMatrix buildCagraGraphForSubset( */ public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) throws IOException { - return writeGraph(graph, vectorIndex, 1); + return writeGraph(graph, vectorIndex, AcceleratedHNSWParams.DEFAULT_GRAPH_THREADS); } - public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int numThreads) + static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int graphThreads) throws IOException { int countOnLevel0 = graph.size(); int numLevels = graph.numLevels(); int[][] offsets = new int[numLevels][]; - - // Each level-0 node's delta/VInt block is independent. Encode those blocks in parallel in fixed - // waves, then concatenate thread buffers in node order. Higher levels are much smaller and stay - // serial. Both paths write the same blocks and offsets in the same order. - // graph.maxConn() scans every layer-0 adjacency row (O(graph size)); compute it once here - // rather than per level/per task below. int maxConn = graph.maxConn(); int[] level0Nodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(0)); offsets[0] = new int[level0Nodes.length]; - if (numThreads > 1 && level0Nodes.length >= PARALLEL_MIN_NODES) { + if (graphThreads > 1 && level0Nodes.length >= GPUBuiltHnswGraph.PARALLEL_MIN_NODES) { writeLevel0Parallel( - graph, vectorIndex, level0Nodes, offsets[0], countOnLevel0, maxConn, numThreads); + graph, vectorIndex, level0Nodes, offsets[0], countOnLevel0, maxConn, graphThreads); } else { writeLevelSerial(graph, vectorIndex, 0, level0Nodes, offsets[0], countOnLevel0, maxConn); } @@ -394,36 +398,32 @@ public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorInde return offsets; } - /** Node count below which parallel level-0 serialization is not worth the overhead. */ - static final int PARALLEL_MIN_NODES = 1 << 16; + /** Absolute node cap for one serialization wave. */ + static final int MAX_SERIALIZATION_WAVE_NODES = 1 << 20; + + /** Maximum worst-case encoded payload buffered by one serialization wave. */ + static final long MAX_SERIALIZED_BYTES_PER_WAVE = 64L << 20; - /** Nodes per wave, bounding the number of nodes buffered independently of dataset size. */ - static final int SERIALIZATION_WAVE_NODES = 1 << 20; + private static final int MAX_VINT_BYTES = 5; - /** Serially encodes a level's nodes into {@code out}, recording per-node byte lengths. */ private static void writeLevelSerial( GPUBuiltHnswGraph graph, IndexOutput out, int level, - int[] sortedNodes, + int[] nodes, int[] offsets, int countOnLevel0, int maxConn) throws IOException { int[] scratch = new int[maxConn * 2]; - int idx = 0; - for (int node : sortedNodes) { + for (int i = 0; i < nodes.length; i++) { long start = out.getFilePointer(); - encodeNode(graph.getNeighbors(level, node), scratch, out, countOnLevel0); - offsets[idx++] = Math.toIntExact(out.getFilePointer() - start); + encodeNode(graph.getNeighbors(level, nodes[i]), scratch, out, countOnLevel0); + offsets[i] = Math.toIntExact(out.getFilePointer() - start); } } - /** - * Encodes level 0 in parallel: within fixed-size waves, threads encode contiguous node - * sub-ranges into per-thread buffers, which are then concatenated to {@code out} in node order - * (identical layout to the serial path). - */ + /** Encodes level zero in bounded waves, then concatenates buffers in node order. */ private static void writeLevel0Parallel( GPUBuiltHnswGraph graph, IndexOutput out, @@ -431,76 +431,68 @@ private static void writeLevel0Parallel( int[] offsets, int countOnLevel0, int maxConn, - int numThreads) + int graphThreads) throws IOException { - // invokeAll joins every task before it returns, including tasks still running when a sibling - // fails, so `buffers` is never concatenated while a worker might still be writing into it. - // It also runs one share of each wave on the calling thread instead of parking it, so the - // pool only has to cover the other ranges. - ExecutorService pool = Executors.newFixedThreadPool(Math.max(1, numThreads - 1)); - try { - TaskExecutor executor = new TaskExecutor(pool); - int n = nodes.length; - for (int waveStart = 0; waveStart < n; ) { - int waveEnd = (int) Math.min(n, (long) waveStart + SERIALIZATION_WAVE_NODES); - int perThread = (waveEnd - waveStart + numThreads - 1) / numThreads; - - ByteBuffersDataOutput[] buffers = new ByteBuffersDataOutput[numThreads]; - List> tasks = new ArrayList<>(numThreads); - for (int t = 0; t < numThreads; t++) { - final int subStart = waveStart + t * perThread; - final int subEnd = Math.min(subStart + perThread, waveEnd); - final int slot = t; - if (subStart >= subEnd) { - continue; - } - tasks.add( - () -> { - ByteBuffersDataOutput buffer = new ByteBuffersDataOutput(); - int[] scratch = new int[maxConn * 2]; - for (int i = subStart; i < subEnd; i++) { - long before = buffer.size(); - encodeNode(graph.getNeighbors(0, nodes[i]), scratch, buffer, countOnLevel0); - offsets[i] = Math.toIntExact(buffer.size() - before); - } - buffers[slot] = buffer; - return null; - }); + int waveNodes = serializationWaveNodes(maxConn); + for (int waveStart = 0; waveStart < nodes.length; ) { + int waveEnd = (int) Math.min(nodes.length, (long) waveStart + waveNodes); + int nodesPerTask = Math.ceilDiv(waveEnd - waveStart, graphThreads); + ByteBuffersDataOutput[] buffers = new ByteBuffersDataOutput[graphThreads]; + List> tasks = new ArrayList<>(graphThreads); + for (int task = 0; task < graphThreads; task++) { + int start = waveStart + task * nodesPerTask; + int end = Math.min(start + nodesPerTask, waveEnd); + int bufferIndex = task; + if (start >= end) { + break; } - executor.invokeAll(tasks); - // Concatenate in thread order (== node order), preserving the serial byte layout. - for (ByteBuffersDataOutput buffer : buffers) { - if (buffer != null) { - buffer.copyTo(out); - } + tasks.add( + () -> { + ByteBuffersDataOutput buffer = new ByteBuffersDataOutput(); + int[] scratch = new int[maxConn * 2]; + for (int i = start; i < end; i++) { + long before = buffer.size(); + encodeNode(graph.getNeighbors(0, nodes[i]), scratch, buffer, countOnLevel0); + offsets[i] = Math.toIntExact(buffer.size() - before); + } + buffers[bufferIndex] = buffer; + return null; + }); + } + GraphWorkExecutor.invokeAll(tasks); + for (ByteBuffersDataOutput buffer : buffers) { + if (buffer != null) { + buffer.copyTo(out); } - waveStart = waveEnd; } - } finally { - pool.shutdown(); + waveStart = waveEnd; } } - /** - * Sorts, delta-encodes and de-duplicates a node's neighbors and writes the block (VInt size + VInt - * deltas) to {@code out}. Shared by the serial and parallel paths so encoding is identical. - */ + static int serializationWaveNodes(int maxConn) { + if (maxConn < 0) { + throw new IllegalArgumentException("maxConn must not be negative"); + } + long maxBytesPerNode = Math.addExact(MAX_VINT_BYTES, (long) maxConn * MAX_VINT_BYTES); + long byteBoundedNodes = Math.max(1, MAX_SERIALIZED_BYTES_PER_WAVE / maxBytesPerNode); + return (int) Math.min(MAX_SERIALIZATION_WAVE_NODES, byteBoundedNodes); + } + private static void encodeNode( NeighborArray neighbors, int[] scratch, DataOutput out, int countOnLevel0) throws IOException { int size = neighbors == null ? 0 : neighbors.size(); int actualSize = 0; if (size > 0) { - int[] nnodes = neighbors.nodes(); - Arrays.sort(nnodes, 0, size); - scratch[0] = nnodes[0]; + int[] nodes = neighbors.nodes(); + Arrays.sort(nodes, 0, size); + scratch[0] = nodes[0]; actualSize = 1; for (int i = 1; i < size; i++) { - assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; - if (nnodes[i - 1] == nnodes[i]) { - continue; + assert nodes[i] < countOnLevel0 : "node too large: " + nodes[i] + ">=" + countOnLevel0; + if (nodes[i - 1] != nodes[i]) { + scratch[actualSize++] = nodes[i] - nodes[i - 1]; } - scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; } } out.writeVInt(actualSize); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index f6f15eef5f..2354e11be8 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -13,11 +13,9 @@ import java.io.IOException; import java.util.ArrayList; import java.util.List; +import java.util.Optional; import java.util.concurrent.Callable; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import java.util.function.Supplier; -import org.apache.lucene.search.TaskExecutor; import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.NeighborArray; @@ -39,6 +37,9 @@ public class GPUBuiltHnswGraph extends HnswGraph { // Layer 0 is special - it contains all nodes private final NeighborArray[] layer0Neighbors; + /** Node count below which parallel graph processing is not worth its overhead. */ + static final int PARALLEL_MIN_NODES = 1 << 16; + private record MaterializedGraph( int numLevels, List layerNodes, @@ -58,15 +59,15 @@ public GPUBuiltHnswGraph( this(size, dimensions, materializeSerial(size, layerNodes, layerAdjacencies)); } - /** Builds a graph while materializing adjacency rows with the requested number of threads. */ - public GPUBuiltHnswGraph( + // Builds a graph while materializing adjacency rows with up to graphThreads threads. + GPUBuiltHnswGraph( int size, int dimensions, List layerNodes, List layerAdjacencies, - int numThreads) + int graphThreads) throws IOException { - this(size, dimensions, materialize(size, layerNodes, layerAdjacencies, numThreads)); + this(size, dimensions, materialize(size, layerNodes, layerAdjacencies, graphThreads)); } private GPUBuiltHnswGraph(int size, int dimensions, MaterializedGraph graph) { @@ -80,107 +81,87 @@ private GPUBuiltHnswGraph(int size, int dimensions, MaterializedGraph graph) { private static MaterializedGraph materializeSerial( int size, List layerNodes, List layerAdjacencies) { - List upperLayerNodes = new ArrayList<>(); - List upperLayerNeighbors = new ArrayList<>(); - NeighborArray[] baseLayerNeighbors = fillNeighborArraySerial(layerAdjacencies.get(0), size); - - for (int level = 1; level < layerAdjacencies.size(); level++) { - int[] nodes = layerNodes.get(level); - upperLayerNodes.add(nodes); - upperLayerNeighbors.add(fillNeighborArraySerial(layerAdjacencies.get(level), nodes.length)); + try { + return materialize(size, layerNodes, layerAdjacencies, 1); + } catch (IOException impossible) { + throw new AssertionError( + "serial graph materialization cannot fail with IOException", impossible); } - return new MaterializedGraph( - layerAdjacencies.size(), upperLayerNodes, baseLayerNeighbors, upperLayerNeighbors); } private static MaterializedGraph materialize( - int size, List layerNodes, List layerAdjacencies, int numThreads) + int size, List layerNodes, List layerAdjacencies, int graphThreads) throws IOException { - if (numThreads <= 1) { - return materializeSerial(size, layerNodes, layerAdjacencies); - } - List upperLayerNodes = new ArrayList<>(); List upperLayerNeighbors = new ArrayList<>(); NeighborArray[] baseLayerNeighbors = - fillNeighborArray(layerAdjacencies.get(0), size, numThreads); + fillNeighborArray(layerAdjacencies.get(0), size, graphThreads); for (int level = 1; level < layerAdjacencies.size(); level++) { int[] nodes = layerNodes.get(level); upperLayerNodes.add(nodes); upperLayerNeighbors.add( - fillNeighborArray(layerAdjacencies.get(level), nodes.length, numThreads)); + fillNeighborArray(layerAdjacencies.get(level), nodes.length, graphThreads)); } return new MaterializedGraph( layerAdjacencies.size(), upperLayerNodes, baseLayerNeighbors, upperLayerNeighbors); } - /** Node count below which parallel materialization is not worth the thread overhead. */ - static final int PARALLEL_MIN_NODES = 1 << 16; - /** - * Maximum temporary native-host copy used to make a device adjacency safe for concurrent reads. - * Larger device matrices retain serial row access instead of risking a full-matrix native-memory - * spike on top of the Java {@link NeighborArray} representation. - */ - static final long MAX_PARALLEL_GRAPH_COPY_BYTES = 4L << 30; - - /** - * Materializes the adjacency matrix into on-heap {@link NeighborArray}s, one per node. - * - *

The serial path reads the adjacency directly (a device matrix's {@code getRow} is safe - * single-threaded). The parallel path cannot: the CAGRA layer-0 adjacency is a device matrix whose - * {@code getRow} uses a shared, stateful buffered reader that is not safe for concurrent access, so - * it is pulled to host once (a single bulk device->host copy) before materializing disjoint node - * ranges concurrently. Host matrices (the upper layers, built via {@link CuVSMatrix#ofArray}) are - * read directly in both paths. + * Materializes an adjacency matrix into heap-backed neighbor arrays. Device matrices are copied + * to host before parallel reads because their row reader is stateful. * * @param adjacency instance of adjacency CuVSMatrix * @param size the number of nodes - * @param numThreads threads to use (1, or fewer than {@value #PARALLEL_MIN_NODES} nodes = serial) + * @param graphThreads maximum threads for this operation, including the caller * @return the NeighborArray */ - private static NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size, int numThreads) + private static NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size, int graphThreads) throws IOException { - NeighborArray[] neighbors = new NeighborArray[size]; - if (numThreads <= 1 - || size < PARALLEL_MIN_NODES - || (adjacency instanceof CuVSDeviceMatrix - && !fitsParallelGraphCopyBudget(adjacency.size(), adjacency.columns()))) { - fillNeighborRange(adjacency, neighbors, 0, size); - return neighbors; + if (graphThreads <= 1 || size < PARALLEL_MIN_NODES) { + return fillNeighborArraySerial(adjacency, size); } + if (adjacency instanceof CuVSDeviceMatrix deviceAdjacency) { - try (CuVSHostMatrix hostCopy = copyToHost(deviceAdjacency)) { - fillNeighborArrayParallel(hostCopy, neighbors, size, numThreads); - } - return neighbors; + return materializeDeviceAdjacency( + deviceAdjacency, + size, + graphThreads, + GraphCopyMemoryBudget.system(), + () -> newHostMatrix(deviceAdjacency)); } - fillNeighborArrayParallel(adjacency, neighbors, size, numThreads); - return neighbors; + + return fillNeighborArrayParallel(adjacency, size, graphThreads); } - private static NeighborArray[] fillNeighborArraySerial(CuVSMatrix adjacency, int size) { + private static NeighborArray[] fillNeighborArraySerial(CuVSMatrix source, int size) { NeighborArray[] neighbors = new NeighborArray[size]; - fillNeighborRange(adjacency, neighbors, 0, size); + fillNeighborRange(source, neighbors, 0, size); return neighbors; } - /** Returns whether an INT32 adjacency can be copied without exceeding the native-host budget. */ - static boolean fitsParallelGraphCopyBudget(long rows, long columns) { - if (rows < 0 || columns < 0) { - return false; + static NeighborArray[] materializeDeviceAdjacency( + CuVSDeviceMatrix source, + int size, + int graphThreads, + GraphCopyMemoryBudget memoryBudget, + Supplier hostCopyFactory) + throws IOException { + Optional reservation = + memoryBudget.tryReserve(source.size(), source.columns()); + if (reservation.isEmpty()) { + return fillNeighborArraySerial(source, size); } - if (rows == 0 || columns == 0) { - return true; + try (GraphCopyMemoryBudget.Reservation ignored = reservation.orElseThrow(); + CuVSHostMatrix hostCopy = copyToHost(source, hostCopyFactory)) { + return fillNeighborArrayParallel(hostCopy, size, graphThreads); } - return rows <= MAX_PARALLEL_GRAPH_COPY_BYTES / Integer.BYTES / columns; } - private static CuVSHostMatrix copyToHost(CuVSDeviceMatrix source) { + private static CuVSHostMatrix newHostMatrix(CuVSDeviceMatrix source) { try (CuVSMatrix.Builder builder = CuVSMatrix.hostBuilder(source.size(), source.columns(), source.dataType())) { - return copyToHost(source, builder::build); + return builder.build(); } } @@ -202,49 +183,40 @@ static CuVSHostMatrix copyToHost( } } - /** - * Materializes disjoint node ranges concurrently. Each thread writes its own slots of {@code - * neighbors} and its own {@link NeighborArray} instances, so no synchronization is needed; {@code - * source} must be a host matrix (stateless {@code getRow}). - */ - private static void fillNeighborArrayParallel( - CuVSMatrix source, NeighborArray[] neighbors, int size, int numThreads) throws IOException { - ExecutorService pool = Executors.newFixedThreadPool(Math.max(1, numThreads - 1)); - try { - int perThread = (size + numThreads - 1) / numThreads; - List> tasks = new ArrayList<>(numThreads); - for (int t = 0; t < numThreads; t++) { - final int start = t * perThread; - final int end = Math.min(start + perThread, size); - if (start >= end) { - break; - } - tasks.add( - () -> { - fillNeighborRange(source, neighbors, start, end); - return null; - }); + private static NeighborArray[] fillNeighborArrayParallel( + CuVSMatrix source, int size, int graphThreads) throws IOException { + NeighborArray[] neighbors = new NeighborArray[size]; + int nodesPerTask = Math.ceilDiv(size, graphThreads); + List> tasks = new ArrayList<>(graphThreads); + for (int task = 0; task < graphThreads; task++) { + int start = task * nodesPerTask; + int end = Math.min(start + nodesPerTask, size); + if (start >= end) { + break; } - new TaskExecutor(pool).invokeAll(tasks); - } finally { - pool.shutdown(); + tasks.add( + () -> { + fillNeighborRange(source, neighbors, start, end); + return null; + }); } + GraphWorkExecutor.invokeAll(tasks); + return neighbors; } - /** Fills {@code neighbors[start, end)} from the adjacency rows. */ private static void fillNeighborRange( CuVSMatrix source, NeighborArray[] neighbors, int start, int end) { for (int i = start; i < end; i++) { RowView rv = source.getRow(i); if (rv != null && rv.size() > 0) { - NeighborArray na = new NeighborArray((int) rv.size(), true); + NeighborArray nodeNeighbors = new NeighborArray((int) rv.size(), true); for (int j = 0; j < rv.size(); j++) { int neighbor = rv.getAsInt(j); if (neighbor >= 0) { - na.addInOrder(neighbor, 1.0f - (j * 0.001f)); + nodeNeighbors.addInOrder(neighbor, 1.0f - (j * 0.001f)); } } - neighbors[i] = na; + neighbors[i] = nodeNeighbors; } else { neighbors[i] = new NeighborArray(0, true); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java new file mode 100644 index 0000000000..0a03e884d8 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java @@ -0,0 +1,137 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.sun.management.OperatingSystemMXBean; +import java.lang.management.ManagementFactory; +import java.util.Objects; +import java.util.Optional; +import org.apache.lucene.util.RamUsageEstimator; +import org.apache.lucene.util.hnsw.NeighborArray; + +/** Coordinates temporary native graph copies across concurrent segment flushes. */ +final class GraphCopyMemoryBudget { + private static final long NEIGHBOR_ARRAY_SHALLOW_BYTES = + RamUsageEstimator.shallowSizeOfInstance(NeighborArray.class); + + private static final GraphCopyMemoryBudget SYSTEM = + new GraphCopyMemoryBudget(GraphCopyMemoryBudget::readSystemMemory); + + private final MemoryProbe memoryProbe; + private long reservedHeadroomBytes; + + GraphCopyMemoryBudget(MemoryProbe memoryProbe) { + this.memoryProbe = Objects.requireNonNull(memoryProbe); + } + + static GraphCopyMemoryBudget system() { + return SYSTEM; + } + + /** + * Tries to reserve enough observed free memory for one graph copy and its materialized graph. + * Callers in the same class loader share reservations. This is cooperative admission control, + * not an operating-system memory guarantee. + */ + synchronized Optional tryReserve(long rows, long columns) { + long requiredHeadroom = requiredHeadroom(rows, columns); + if (requiredHeadroom < 0) { + return Optional.empty(); + } + + MemorySnapshot memory; + try { + memory = memoryProbe.read(); + } catch (RuntimeException unavailable) { + return Optional.empty(); + } + if (memory == null + || memory.totalBytes() <= 0 + || memory.freeBytes() < 0 + || memory.freeBytes() > memory.totalBytes()) { + return Optional.empty(); + } + + if (reservedHeadroomBytes > memory.freeBytes() + || requiredHeadroom > memory.freeBytes() - reservedHeadroomBytes) { + return Optional.empty(); + } + reservedHeadroomBytes += requiredHeadroom; + return Optional.of(new Reservation(this, requiredHeadroom)); + } + + /** + * Estimates peak allocation from the actual matrix shape and current JVM object layout. Besides + * the native host copy and materialized Lucene graph, one adjacency-sized allowance protects + * against allocation races and estimation error while the copy is in flight. + */ + static long requiredHeadroom(long rows, long columns) { + if (rows <= 0 || rows > Integer.MAX_VALUE || columns <= 0 || columns > Integer.MAX_VALUE) { + return -1; + } + try { + long adjacencyBytes = Math.multiplyExact(Math.multiplyExact(rows, columns), Integer.BYTES); + long neighborReferences = arraySize(rows, RamUsageEstimator.NUM_BYTES_OBJECT_REF); + long nodeIds = arraySize(columns, Integer.BYTES); + long scores = arraySize(columns, Float.BYTES); + long bytesPerNode = Math.addExact(NEIGHBOR_ARRAY_SHALLOW_BYTES, nodeIds); + bytesPerNode = Math.addExact(bytesPerNode, scores); + long luceneGraphBytes = + Math.addExact(neighborReferences, Math.multiplyExact(rows, bytesPerNode)); + return Math.addExact(Math.multiplyExact(adjacencyBytes, 2), luceneGraphBytes); + } catch (ArithmeticException overflow) { + return -1; + } + } + + private static long arraySize(long length, int bytesPerElement) { + long unaligned = + Math.addExact( + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER, Math.multiplyExact(length, bytesPerElement)); + long alignment = RamUsageEstimator.NUM_BYTES_OBJECT_ALIGNMENT; + long remainder = unaligned % alignment; + return remainder == 0 ? unaligned : Math.addExact(unaligned, alignment - remainder); + } + + private synchronized void release(Reservation reservation) { + if (reservation.released) { + return; + } + reservedHeadroomBytes -= reservation.headroomBytes; + reservation.released = true; + } + + private static MemorySnapshot readSystemMemory() { + java.lang.management.OperatingSystemMXBean platformBean = + ManagementFactory.getOperatingSystemMXBean(); + if (platformBean instanceof OperatingSystemMXBean osBean) { + return new MemorySnapshot(osBean.getTotalMemorySize(), osBean.getFreeMemorySize()); + } + return null; + } + + @FunctionalInterface + interface MemoryProbe { + MemorySnapshot read(); + } + + record MemorySnapshot(long totalBytes, long freeBytes) {} + + static final class Reservation implements AutoCloseable { + private final GraphCopyMemoryBudget budget; + private final long headroomBytes; + private boolean released; + + private Reservation(GraphCopyMemoryBudget budget, long headroomBytes) { + this.budget = budget; + this.headroomBytes = headroomBytes; + } + + @Override + public void close() { + budget.release(this); + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphWorkExecutor.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphWorkExecutor.java new file mode 100644 index 0000000000..cb715f1fd8 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphWorkExecutor.java @@ -0,0 +1,167 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.FutureTask; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.SynchronousQueue; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.ThreadInterruptedException; + +/** Shares a bounded set of helper threads across accelerated-HNSW graph operations. */ +final class GraphWorkExecutor { + static final String THREAD_NAME_PREFIX = "cuvs-hnsw-graph-worker-"; + + private static final int MAX_WORKERS = + defaultMaxWorkers(Runtime.getRuntime().availableProcessors()); + private static final long KEEP_ALIVE_SECONDS = 1; + private static final AtomicInteger NEXT_THREAD_ID = new AtomicInteger(); + private static final ThreadPoolExecutor EXECUTOR = + newExecutor(MAX_WORKERS, KEEP_ALIVE_SECONDS, TimeUnit.SECONDS); + + private GraphWorkExecutor() {} + + static void invokeAll(List> tasks) throws IOException { + invokeAll(EXECUTOR, tasks); + } + + /** + * Runs one task on the calling thread and offers the others to {@code executor}. Direct handoff + * and caller execution apply backpressure without retaining graph work in a queue. + */ + static void invokeAll(Executor executor, List> tasks) throws IOException { + Objects.requireNonNull(executor, "executor"); + Objects.requireNonNull(tasks, "tasks"); + if (tasks.isEmpty()) { + return; + } + AtomicBoolean taskFailed = new AtomicBoolean(); + List> preparedTasks = new ArrayList<>(tasks.size()); + for (Callable task : tasks) { + preparedTasks.add( + new FutureTask<>(stopAfterFailure(Objects.requireNonNull(task, "task"), taskFailed))); + } + + List> acceptedTasks = new ArrayList<>(preparedTasks.size()); + Throwable failure = null; + for (int i = 0; i < preparedTasks.size(); i++) { + FutureTask task = preparedTasks.get(i); + if (i == preparedTasks.size() - 1) { + acceptedTasks.add(task); + task.run(); + break; + } + try { + executor.execute(task); + acceptedTasks.add(task); + } catch (RejectedExecutionException rejected) { + acceptedTasks.add(task); + task.run(); + } catch (RuntimeException | Error submissionFailure) { + taskFailed.set(true); + failure = submissionFailure; + break; + } + } + + failure = collectFailures(acceptedTasks, failure); + if (failure != null) { + throw IOUtils.rethrowAlways(failure); + } + } + + static ThreadPoolExecutor newExecutor(int maxWorkers, long keepAlive, TimeUnit unit) { + if (maxWorkers < 1) { + throw new IllegalArgumentException("maxWorkers must be positive"); + } + if (keepAlive <= 0) { + throw new IllegalArgumentException("keepAlive must be positive"); + } + Objects.requireNonNull(unit, "unit"); + + return new ThreadPoolExecutor( + 0, + maxWorkers, + keepAlive, + unit, + new SynchronousQueue<>(), + workerThreadFactory(), + new ThreadPoolExecutor.AbortPolicy()); + } + + static int defaultMaxWorkers(int availableProcessors) { + return availableProcessors <= 1 ? 1 : availableProcessors - 1; + } + + private static Callable stopAfterFailure(Callable task, AtomicBoolean taskFailed) { + return () -> { + if (taskFailed.get()) { + return null; + } + try { + return task.call(); + } catch (Exception | Error failure) { + taskFailed.set(true); + throw failure; + } + }; + } + + private static ThreadFactory workerThreadFactory() { + return task -> { + Thread worker = + new Thread(null, task, THREAD_NAME_PREFIX + NEXT_THREAD_ID.incrementAndGet(), 0, false); + worker.setDaemon(true); + worker.setPriority(Thread.NORM_PRIORITY); + worker.setContextClassLoader(null); + return worker; + }; + } + + /** Waits for all accepted work before returning, even after failure or interruption. */ + private static Throwable collectFailures(List> tasks, Throwable failure) { + boolean interrupted = false; + for (FutureTask task : tasks) { + while (true) { + try { + task.get(); + break; + } catch (InterruptedException interruption) { + interrupted = true; + failure = addFailure(failure, new ThreadInterruptedException(interruption)); + } catch (ExecutionException taskFailure) { + failure = addFailure(failure, taskFailure.getCause()); + break; + } + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + return failure; + } + + private static Throwable addFailure(Throwable failure, Throwable next) { + if (failure == null) { + return next; + } + if (failure != next) { + failure.addSuppressed(next); + } + return failure; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 8a27faf957..f113f5e198 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -186,10 +186,10 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.NONE, - acceleratedHNSWParams.getWriterThreads()); + acceleratedHNSWParams.getGraphThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getGraphThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -275,7 +275,7 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getGraphThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 993fb8a58b..3cf50be3b7 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -184,12 +184,11 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.BINARY, - acceleratedHNSWParams.getWriterThreads()); + acceleratedHNSWParams.getGraphThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - // Write the graph to the vector index int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getGraphThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -284,7 +283,7 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { // Write the graph to the vector index int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getGraphThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 0291b928ad..77370fb26b 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -211,12 +211,11 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.SCALAR, - acceleratedHNSWParams.getWriterThreads()); + acceleratedHNSWParams.getGraphThreads()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - // Write the graph to the vector index int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getGraphThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -310,7 +309,7 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getGraphThreads()); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/IntGraphTestMatrix.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/IntGraphTestMatrix.java new file mode 100644 index 0000000000..73ef7c5a78 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/IntGraphTestMatrix.java @@ -0,0 +1,238 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.RowView; +import java.util.Random; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +/** Minimal in-memory INT matrix used by graph materialization and serialization tests. */ +class IntGraphTestMatrix implements CuVSMatrix { + private final int[][] rows; + + IntGraphTestMatrix(int[][] rows) { + this.rows = rows; + } + + static IntGraphTestMatrix random(int rowCount, int columnCount, long seed) { + return new IntGraphTestMatrix(randomRows(rowCount, columnCount, seed)); + } + + static int[][] randomRows(int rowCount, int columnCount, long seed) { + Random random = new Random(seed); + int[][] rows = new int[rowCount][columnCount]; + for (int[] row : rows) { + for (int column = 0; column < row.length; column++) { + row[column] = random.nextInt(rowCount); + } + } + return rows; + } + + @Override + public long size() { + return rows.length; + } + + @Override + public long columns() { + return rows.length == 0 ? 0 : rows[0].length; + } + + @Override + public DataType dataType() { + return DataType.INT; + } + + @Override + public RowView getRow(long row) { + return new IntRow(rows[Math.toIntExact(row)]); + } + + @Override + public void toArray(int[][] target) { + for (int row = 0; row < rows.length; row++) { + System.arraycopy(rows[row], 0, target[row], 0, rows[row].length); + } + } + + @Override + public void toArray(float[][] target) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(byte[][] target) { + throw new UnsupportedOperationException(); + } + + @Override + public void toHost(CuVSHostMatrix target) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSHostMatrix toHost() { + throw new UnsupportedOperationException(); + } + + @Override + public void toDevice(CuVSDeviceMatrix target, CuVSResources resources) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSDeviceMatrix toDevice(CuVSResources resources) { + throw new UnsupportedOperationException(); + } + + @Override + public void close() {} + + /** Reports a synthetic device shape and fails if an unapproved host copy is attempted. */ + static class DeviceMatrix extends IntGraphTestMatrix implements CuVSDeviceMatrix { + private final long reportedColumns; + + DeviceMatrix(int[][] rows, long reportedColumns) { + super(rows); + this.reportedColumns = reportedColumns; + } + + @Override + public long columns() { + return reportedColumns; + } + + @Override + public void toHost(CuVSHostMatrix target) { + throw new AssertionError("rejected device adjacency must not be copied to host"); + } + + @Override + public CuVSHostMatrix toHost() { + throw new AssertionError("rejected device adjacency must not be copied to host"); + } + } + + static final class TrackingHostMatrix extends IntGraphTestMatrix implements CuVSHostMatrix { + private final AtomicInteger closeCount; + private final RuntimeException closeFailure; + private final ParallelExecutionProbe executionProbe; + + TrackingHostMatrix(AtomicInteger closeCount, RuntimeException closeFailure) { + this(new int[][] {{0}}, closeCount, closeFailure, null); + } + + TrackingHostMatrix( + int[][] rows, + AtomicInteger closeCount, + RuntimeException closeFailure, + ParallelExecutionProbe executionProbe) { + super(rows); + this.closeCount = closeCount; + this.closeFailure = closeFailure; + this.executionProbe = executionProbe; + } + + @Override + public int get(int row, int column) { + return getRow(row).getAsInt(column); + } + + @Override + public RowView getRow(long row) { + if (executionProbe != null) { + executionProbe.recordExecution(); + } + return super.getRow(row); + } + + @Override + public void close() { + closeCount.incrementAndGet(); + if (closeFailure != null) { + throw closeFailure; + } + } + } + + /** Holds the first operation until a second execution context reaches the same work. */ + static final class ParallelExecutionProbe { + private static final long TIMEOUT_SECONDS = 10; + + private final Set threads = ConcurrentHashMap.newKeySet(); + private final CountDownLatch parallelExecution = new CountDownLatch(1); + + void recordExecution() { + threads.add(Thread.currentThread()); + if (threads.size() > 1) { + parallelExecution.countDown(); + } + try { + if (!parallelExecution.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)) { + throw new AssertionError("work never reached a second execution context"); + } + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new AssertionError("interrupted while observing parallel work", interrupted); + } + } + + int threadCount() { + return threads.size(); + } + } + + private static final class IntRow implements RowView { + private final int[] values; + + private IntRow(int[] values) { + this.values = values; + } + + @Override + public long size() { + return values.length; + } + + @Override + public int getAsInt(long index) { + return values[Math.toIntExact(index)]; + } + + @Override + public float getAsFloat(long index) { + throw new UnsupportedOperationException(); + } + + @Override + public byte getAsByte(long index) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(int[] target) { + System.arraycopy(values, 0, target, 0, values.length); + } + + @Override + public void toArray(float[] target) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(byte[] target) { + throw new UnsupportedOperationException(); + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java index e5818c1ba2..6794ad1f89 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java @@ -9,6 +9,7 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_CAGRA_GRAPH_BUILD_ALGO; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_CUVS_DISTANCE_TYPE; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_GRAPH_DEGREE; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_GRAPH_THREADS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_HNSW_LAYERS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_INT_GRAPH_DEGREE; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_MAX_CONN; @@ -18,6 +19,7 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_WRITER_THREADS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_BEAM_WIDTH; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_GRAPH_THREADS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_HNSW_LAYERS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_INT_GRAPH_DEG; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_MAX_CONN; @@ -26,6 +28,7 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_WRITER_THREADS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_BEAM_WIDTH; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_GRAPH_THREADS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_HNSW_LAYERS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_INT_GRAPH_DEG; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_MAX_CONN; @@ -55,6 +58,7 @@ public void testAcceleratedHNSWParamsDefaultValues() { AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); assertEquals(DEFAULT_BEAM_WIDTH, params.getBeamWidth()); assertEquals(DEFAULT_GRAPH_DEGREE, params.getGraphdegree()); + assertEquals(DEFAULT_GRAPH_THREADS, params.getGraphThreads()); assertEquals(DEFAULT_HNSW_LAYERS, params.getHnswLayers()); assertEquals(DEFAULT_INT_GRAPH_DEGREE, params.getIntermediateGraphDegree()); assertEquals(DEFAULT_MAX_CONN, params.getMaxConn()); @@ -208,6 +212,28 @@ public void testAcceleratedHNSWParamsInvalidWriterThreads() { } } + @Test + public void testAcceleratedHNSWParamsInvalidGraphThreads() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, MIN_GRAPH_THREADS), + random.nextInt(MAX_GRAPH_THREADS + 1, Integer.MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withGraphThreads(v).build()); + } + } + + @Test + public void testWriterAndGraphThreadsAreIndependent() { + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withWriterThreads(7).withGraphThreads(3).build(); + + assertEquals(7, params.getWriterThreads()); + assertEquals(3, params.getGraphThreads()); + } + @Test public void testAcceleratedHNSWParamsInvalidNumMergeWorkers() { for (int v : diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java index f559ea3808..c4319e6f46 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java @@ -184,6 +184,7 @@ public void testHnswHeuristicDelegatesToCuVS() { .withMaxConn(16) .withBeamWidth(100) .withWriterThreads(7) + .withGraphThreads(3) .build(); CagraIndexParams cagraParams = CagraIndexParamsFactory.create(params, 10_000, 128); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java new file mode 100644 index 0000000000..6d0d309d3f --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java @@ -0,0 +1,152 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.RamUsageEstimator; +import org.apache.lucene.util.hnsw.NeighborArray; +import org.junit.Test; + +/** Behavioral specifications for native graph-copy admission control. */ +public class TestGraphCopyMemoryBudget extends LuceneTestCase { + private static final long TIMEOUT_SECONDS = 10; + + @Test + public void reservationsFollowEstimatedPeakAcrossSupportedDegrees() { + int rows = 100; + for (int degree : new int[] {1, 32, 512}) { + long required = GraphCopyMemoryBudget.requiredHeadroom(rows, degree); + long adjacencyBytes = (long) rows * degree * Integer.BYTES; + long expected = + 2 * adjacencyBytes + + RamUsageEstimator.shallowSizeOf(new NeighborArray[rows]) + + rows + * (RamUsageEstimator.shallowSizeOfInstance(NeighborArray.class) + + RamUsageEstimator.sizeOf(new int[degree]) + + RamUsageEstimator.sizeOf(new float[degree])); + assertEquals(expected, required); + assertTrue("object layout must be included", required > 4 * adjacencyBytes); + + GraphCopyMemoryBudget exactBudget = budgetWith(required, required); + try (GraphCopyMemoryBudget.Reservation ignored = reserve(exactBudget, rows, degree)) { + assertTrue(exactBudget.tryReserve(1, 1).isEmpty()); + } + + GraphCopyMemoryBudget insufficientBudget = budgetWith(required, required - 1); + assertTrue(insufficientBudget.tryReserve(rows, degree).isEmpty()); + } + } + + @Test + public void concurrentReservationsCannotExceedSharedHeadroom() throws Exception { + long rows = 100; + long degree = 16; + long reservationBytes = GraphCopyMemoryBudget.requiredHeadroom(rows, degree); + GraphCopyMemoryBudget budget = budgetWith(2 * reservationBytes, 2 * reservationBytes); + int callers = 8; + ExecutorService executor = Executors.newFixedThreadPool(callers); + CountDownLatch start = new CountDownLatch(1); + CountDownLatch attempted = new CountDownLatch(callers); + CountDownLatch release = new CountDownLatch(1); + AtomicInteger granted = new AtomicInteger(); + List> futures = new ArrayList<>(); + try { + for (int i = 0; i < callers; i++) { + futures.add( + executor.submit( + () -> { + assertTrue(start.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + Optional reservation = + budget.tryReserve(rows, degree); + reservation.ifPresent(ignored -> granted.incrementAndGet()); + attempted.countDown(); + if (reservation.isPresent()) { + assertTrue(release.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + reservation.orElseThrow().close(); + } + return null; + })); + } + + start.countDown(); + assertTrue(attempted.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + assertEquals(2, granted.get()); + release.countDown(); + for (Future future : futures) { + future.get(TIMEOUT_SECONDS, TimeUnit.SECONDS); + } + } finally { + release.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } + } + + @Test + public void reservationIsReleasedOnFailureAndCloseIsIdempotent() { + long required = GraphCopyMemoryBudget.requiredHeadroom(100, 16); + GraphCopyMemoryBudget budget = budgetWith(required, required); + GraphCopyMemoryBudget.Reservation failedOperation = reserve(budget, 100, 16); + RuntimeException expected = new RuntimeException("expected"); + + RuntimeException actual = + assertThrows( + RuntimeException.class, + () -> { + try (failedOperation) { + throw expected; + } + }); + assertSame(expected, actual); + failedOperation.close(); + + try (GraphCopyMemoryBudget.Reservation replacement = reserve(budget, 100, 16)) { + assertTrue(budget.tryReserve(1, 1).isEmpty()); + } + } + + @Test + public void invalidOrUnavailableMemoryInformationFailsClosed() { + assertRejected(() -> null); + assertRejected(() -> new GraphCopyMemoryBudget.MemorySnapshot(0, 0)); + assertRejected(() -> new GraphCopyMemoryBudget.MemorySnapshot(1_000, -1)); + assertRejected(() -> new GraphCopyMemoryBudget.MemorySnapshot(1_000, 1_001)); + assertRejected( + () -> { + throw new UnsupportedOperationException("unavailable"); + }); + + GraphCopyMemoryBudget budget = + budgetWith(/* totalBytes= */ Long.MAX_VALUE, /* freeBytes= */ Long.MAX_VALUE); + assertTrue(budget.tryReserve(Integer.MAX_VALUE, Integer.MAX_VALUE).isEmpty()); + assertTrue(budget.tryReserve(0, 1).isEmpty()); + assertTrue(budget.tryReserve(1, 0).isEmpty()); + assertTrue(budget.tryReserve(-1, 1).isEmpty()); + assertTrue(budget.tryReserve(1, -1).isEmpty()); + } + + private static GraphCopyMemoryBudget budgetWith(long totalBytes, long freeBytes) { + return new GraphCopyMemoryBudget( + () -> new GraphCopyMemoryBudget.MemorySnapshot(totalBytes, freeBytes)); + } + + private static GraphCopyMemoryBudget.Reservation reserve( + GraphCopyMemoryBudget budget, long rows, long columns) { + return budget.tryReserve(rows, columns).orElseThrow(); + } + + private static void assertRejected(GraphCopyMemoryBudget.MemoryProbe probe) { + assertTrue(new GraphCopyMemoryBudget(probe).tryReserve(1, 1).isEmpty()); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphThreadsPersistedIndex.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphThreadsPersistedIndex.java new file mode 100644 index 0000000000..cf980d439b --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphThreadsPersistedIndex.java @@ -0,0 +1,127 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import java.util.Random; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.HnswGraphProvider; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestGraphThreadsPersistedIndex extends LuceneTestCase { + + private static final String VECTOR_FIELD = "vector"; + private static final int VECTOR_COUNT = GPUBuiltHnswGraph.PARALLEL_MIN_NODES + 1; + private static final int DIMENSIONS = 32; + + @Test + public void testGraphThreadsRoundTripAboveThreshold() throws Exception { + assumeTrue("cuVS not supported", isSupported()); + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder() + .withWriterThreads(1) + .withGraphThreads(4) + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withIntermediateGraphDegree(32) + .withGraphDegree(16) + .withHNSWLayer(1) + .build(); + Codec codec = new Lucene101AcceleratedHNSWCodec(params); + Random random = new Random(0x2594L); + + try (Directory directory = newDirectory()) { + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setUseCompoundFile(false) + .setMaxBufferedDocs(VECTOR_COUNT + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int id = 0; id < VECTOR_COUNT; id++) { + float[] vector = new float[DIMENSIONS]; + for (int dimension = 0; dimension < DIMENSIONS; dimension++) { + vector[dimension] = random.nextFloat(); + } + Document document = new Document(); + document.add(new StringField("id", Integer.toString(id), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, vector, EUCLIDEAN)); + writer.addDocument(document); + } + } + + TestUtil.checkIndex(directory); + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals(1, reader.leaves().size()); + assertEquals(VECTOR_COUNT, reader.numDocs()); + LeafReader leaf = getOnlyLeafReader(reader); + HnswGraph graph = graphOf(leaf); + assertEquals(VECTOR_COUNT, graph.size()); + int arcs = 0; + HnswGraph.NodesIterator nodes = graph.getNodesOnLevel(0); + while (nodes.hasNext()) { + int node = nodes.nextInt(); + graph.seek(0, node); + for (int neighbor = graph.nextNeighbor(); + neighbor != NO_MORE_DOCS; + neighbor = graph.nextNeighbor()) { + assertTrue(neighbor >= 0); + assertTrue(neighbor < VECTOR_COUNT); + arcs++; + } + } + assertTrue("persisted graph contains no arcs", arcs > 0); + + int queryNode = graph.entryNode(); + assertTrue(queryNode >= 0); + assertTrue(queryNode < VECTOR_COUNT); + FloatVectorValues values = leaf.getFloatVectorValues(VECTOR_FIELD); + assertNotNull(values); + float[] query = values.vectorValue(queryNode).clone(); + int queryDoc = values.ordToDoc(queryNode); + String queryId = leaf.storedFields().document(queryDoc).get("id"); + + IndexSearcher searcher = new IndexSearcher(reader); + var hits = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, query, 10), 10); + assertEquals(10, hits.scoreDocs.length); + boolean foundQueryNode = false; + for (var hit : hits.scoreDocs) { + foundQueryNode |= queryId.equals(searcher.storedFields().document(hit.doc).get("id")); + } + assertTrue("the entry-node vector must be returned for its own query", foundQueryNode); + } + } + } + + private static HnswGraph graphOf(LeafReader leaf) throws Exception { + KnnVectorsReader reader = ((CodecReader) leaf).getVectorReader(); + if (reader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + reader = fieldsReader.getFieldReader(VECTOR_FIELD); + } + return ((HnswGraphProvider) reader).getGraph(VECTOR_FIELD); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphWorkExecutor.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphWorkExecutor.java new file mode 100644 index 0000000000..9f8cb1c75d --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphWorkExecutor.java @@ -0,0 +1,455 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicIntegerArray; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.ThreadInterruptedException; +import org.junit.Test; + +/** Behavioral specifications for the shared accelerated-HNSW graph-work scheduler. */ +public class TestGraphWorkExecutor extends LuceneTestCase { + private static final long TIMEOUT_SECONDS = 10; + + @Test + public void defaultWorkerCapReservesTheCallingThread() { + assertEquals(1, GraphWorkExecutor.defaultMaxWorkers(1)); + assertEquals(1, GraphWorkExecutor.defaultMaxWorkers(2)); + assertEquals(7, GraphWorkExecutor.defaultMaxWorkers(8)); + } + + @Test + public void saturationRunsWorkOnTheCallerWithoutLosingIt() throws Exception { + ThreadPoolExecutor executor = newExecutor(1); + CountDownLatch helperStarted = new CountDownLatch(1); + CountDownLatch releaseHelper = new CountDownLatch(1); + try { + executor.execute( + () -> { + helperStarted.countDown(); + awaitUninterruptibly(releaseHelper); + }); + assertTrue(helperStarted.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + + AtomicIntegerArray calls = new AtomicIntegerArray(32); + Set taskThreads = ConcurrentHashMap.newKeySet(); + GraphWorkExecutor.invokeAll(executor, countedTasks(calls, taskThreads)); + + assertEquals(Set.of(Thread.currentThread()), taskThreads); + assertEachCalledOnce(calls); + } finally { + releaseHelper.countDown(); + shutdown(executor); + } + } + + @Test + public void concurrentCallersUseTheProductionSharedWorkerBound() throws Exception { + int workerLimit = + GraphWorkExecutor.defaultMaxWorkers(Runtime.getRuntime().availableProcessors()); + int callers = Math.min(4, workerLimit + 1); + int tasksPerCaller = Math.ceilDiv(workerLimit + 1, callers) + 1; + ConcurrentRun run = runConcurrentWork(callers, tasksPerCaller, workerLimit); + + assertEachCalledOnce(run.calls); + assertTrue(run.maxActiveHelpers > 0); + assertTrue(run.maxActiveHelpers <= workerLimit); + } + + @Test + public void failuresWaitForStartedWorkPreserveCausesAndSkipPendingWork() throws Exception { + ThreadPoolExecutor graphExecutor = newExecutor(1); + ExecutorService callerExecutor = Executors.newSingleThreadExecutor(); + CountDownLatch helperStarted = new CountDownLatch(1); + CountDownLatch releaseHelper = new CountDownLatch(1); + CountDownLatch callerTaskFailed = new CountDownLatch(1); + AtomicInteger pendingCalls = new AtomicInteger(); + IOException helperFailure = new IOException("helper failed"); + IllegalStateException callerFailure = new IllegalStateException("caller failed"); + try { + Future invocation = + invokeAsync( + callerExecutor, + () -> + GraphWorkExecutor.invokeAll( + graphExecutor, + List.of( + () -> { + helperStarted.countDown(); + awaitUninterruptibly(releaseHelper); + throw helperFailure; + }, + () -> { + assertTrue(helperStarted.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + callerTaskFailed.countDown(); + throw callerFailure; + }, + () -> { + pendingCalls.incrementAndGet(); + return null; + }))); + + assertTrue(callerTaskFailed.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + assertFalse(invocation.isDone()); + releaseHelper.countDown(); + + Throwable thrown = invocation.get(TIMEOUT_SECONDS, TimeUnit.SECONDS); + assertTrue(thrown == helperFailure || thrown == callerFailure); + Throwable otherFailure = thrown == helperFailure ? callerFailure : helperFailure; + assertArrayEquals(new Throwable[] {otherFailure}, thrown.getSuppressed()); + assertEquals(0, pendingCalls.get()); + } finally { + releaseHelper.countDown(); + shutdown(callerExecutor); + shutdown(graphExecutor); + } + } + + @Test + public void interruptedCallerWaitsForStartedWorkAndRestoresInterrupt() throws Exception { + ThreadPoolExecutor executor = newExecutor(1); + BlockingTask helper = new BlockingTask(); + CountDownLatch callerShareFinished = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + AtomicBoolean interruptedOnExit = new AtomicBoolean(); + try { + Thread caller = + new Thread( + () -> { + try { + GraphWorkExecutor.invokeAll( + executor, + List.of( + helper, + () -> { + helper.awaitStarted(); + callerShareFinished.countDown(); + return null; + })); + } catch (Throwable thrown) { + failure.set(thrown); + } finally { + interruptedOnExit.set(Thread.currentThread().isInterrupted()); + } + }, + "graph-work-interrupted-caller"); + caller.start(); + + assertTrue(callerShareFinished.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + awaitThreadWaiting(caller); + caller.interrupt(); + awaitInterruptConsumed(caller); + assertTrue("caller returned while accepted work was still running", caller.isAlive()); + + helper.release(); + caller.join(TimeUnit.SECONDS.toMillis(TIMEOUT_SECONDS)); + assertFalse("caller did not finish", caller.isAlive()); + assertTrue(failure.get() instanceof ThreadInterruptedException); + assertTrue(interruptedOnExit.get()); + } finally { + helper.release(); + shutdown(executor); + } + } + + @Test + public void submissionFailureWaitsForAcceptedWork() throws Exception { + ThreadPoolExecutor delegate = newExecutor(1); + ExecutorService callerExecutor = Executors.newSingleThreadExecutor(); + BlockingTask acceptedTask = new BlockingTask(); + IllegalStateException expected = new IllegalStateException("submission failed"); + CountDownLatch submissionFailed = new CountDownLatch(1); + AtomicInteger submissions = new AtomicInteger(); + Executor failsAfterFirstSubmission = + task -> { + if (submissions.incrementAndGet() == 1) { + delegate.execute(task); + acceptedTask.awaitStartedUninterruptibly(); + } else { + submissionFailed.countDown(); + throw expected; + } + }; + try { + Future invocation = + invokeAsync( + callerExecutor, + () -> + GraphWorkExecutor.invokeAll( + failsAfterFirstSubmission, List.of(acceptedTask, () -> null, () -> null))); + + assertTrue(submissionFailed.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + assertFalse(invocation.isDone()); + acceptedTask.release(); + + assertSame(expected, invocation.get(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } finally { + acceptedTask.release(); + shutdown(callerExecutor); + shutdown(delegate); + } + } + + @Test + public void workersAreIsolatedDaemonsThatExpireWhenIdle() throws Exception { + ThreadPoolExecutor executor = GraphWorkExecutor.newExecutor(1, 25, TimeUnit.MILLISECONDS); + InheritableThreadLocal callerState = new InheritableThreadLocal<>(); + ClassLoader originalContextLoader = Thread.currentThread().getContextClassLoader(); + ClassLoader callerContextLoader = new ClassLoader(originalContextLoader) {}; + AtomicReference observation = new AtomicReference<>(); + CountDownLatch workerObserved = new CountDownLatch(1); + try { + callerState.set("caller-state"); + Thread.currentThread().setContextClassLoader(callerContextLoader); + Callable observeWorker = + () -> { + if (isGraphWorker()) { + Thread worker = Thread.currentThread(); + observation.set( + new WorkerObservation(worker, callerState.get(), worker.getContextClassLoader())); + workerObserved.countDown(); + } else { + assertTrue(workerObserved.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } + return null; + }; + + GraphWorkExecutor.invokeAll(executor, List.of(observeWorker, observeWorker)); + + WorkerObservation first = observation.get(); + assertNotNull(first); + assertTrue(first.thread.isDaemon()); + assertTrue(first.thread.getName().startsWith(GraphWorkExecutor.THREAD_NAME_PREFIX)); + assertNull(first.inheritedState); + assertNull(first.contextClassLoader); + + awaitPoolSize(executor, 0); + Thread replacement = invokeAndCaptureWorker(executor); + assertNotSame(first.thread, replacement); + } finally { + Thread.currentThread().setContextClassLoader(originalContextLoader); + callerState.remove(); + shutdown(executor); + } + } + + private static ConcurrentRun runConcurrentWork(int callers, int tasksPerCaller, int workerLimit) + throws Exception { + ExecutorService callerExecutor = Executors.newFixedThreadPool(callers); + AtomicIntegerArray calls = new AtomicIntegerArray(callers * tasksPerCaller); + AtomicInteger activeHelpers = new AtomicInteger(); + AtomicInteger maxActiveHelpers = new AtomicInteger(); + CountDownLatch start = new CountDownLatch(1); + CountDownLatch releaseHelpers = new CountDownLatch(1); + CountDownLatch sharedCapacityReached = new CountDownLatch(workerLimit); + List> invocations = new ArrayList<>(); + try { + for (int caller = 0; caller < callers; caller++) { + int invocationId = caller; + invocations.add( + callerExecutor.submit( + () -> { + assertTrue(start.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + List> tasks = new ArrayList<>(tasksPerCaller); + for (int task = 0; task < tasksPerCaller; task++) { + int taskId = invocationId * tasksPerCaller + task; + tasks.add( + () -> { + boolean helper = isGraphWorker(); + if (helper) { + int active = activeHelpers.incrementAndGet(); + maxActiveHelpers.accumulateAndGet(active, Math::max); + sharedCapacityReached.countDown(); + } + try { + calls.incrementAndGet(taskId); + if (helper) { + assertTrue(releaseHelpers.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } + } finally { + if (helper) { + activeHelpers.decrementAndGet(); + } + } + return null; + }); + } + GraphWorkExecutor.invokeAll(tasks); + return null; + })); + } + start.countDown(); + assertTrue(sharedCapacityReached.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + releaseHelpers.countDown(); + awaitAll(invocations); + return new ConcurrentRun(calls, maxActiveHelpers.get()); + } finally { + releaseHelpers.countDown(); + shutdown(callerExecutor); + } + } + + private static List> countedTasks( + AtomicIntegerArray calls, Set taskThreads) { + List> tasks = new ArrayList<>(calls.length()); + for (int task = 0; task < calls.length(); task++) { + int taskId = task; + tasks.add( + () -> { + calls.incrementAndGet(taskId); + taskThreads.add(Thread.currentThread()); + return null; + }); + } + return tasks; + } + + private static Future invokeAsync( + ExecutorService executor, ThrowingAction invocation) { + return executor.submit( + () -> { + try { + invocation.run(); + return null; + } catch (Throwable failure) { + return failure; + } + }); + } + + private static Thread invokeAndCaptureWorker(ThreadPoolExecutor executor) throws IOException { + AtomicReference worker = new AtomicReference<>(); + CountDownLatch observed = new CountDownLatch(1); + Callable capture = + () -> { + if (isGraphWorker()) { + worker.set(Thread.currentThread()); + observed.countDown(); + } else { + assertTrue(observed.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } + return null; + }; + GraphWorkExecutor.invokeAll(executor, List.of(capture, capture)); + assertNotNull(worker.get()); + return worker.get(); + } + + private static boolean isGraphWorker() { + return Thread.currentThread().getName().startsWith(GraphWorkExecutor.THREAD_NAME_PREFIX); + } + + private static ThreadPoolExecutor newExecutor(int maxWorkers) { + return GraphWorkExecutor.newExecutor(maxWorkers, 1, TimeUnit.MINUTES); + } + + private static void assertEachCalledOnce(AtomicIntegerArray calls) { + for (int task = 0; task < calls.length(); task++) { + assertEquals("task " + task, 1, calls.get(task)); + } + } + + private static void awaitAll(List> futures) throws Exception { + for (Future future : futures) { + future.get(TIMEOUT_SECONDS, TimeUnit.SECONDS); + } + } + + private static void awaitPoolSize(ThreadPoolExecutor executor, int expected) + throws InterruptedException { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TIMEOUT_SECONDS); + while (executor.getPoolSize() != expected && System.nanoTime() < deadline) { + Thread.sleep(10); + } + assertEquals("pool size did not reach " + expected, expected, executor.getPoolSize()); + } + + private static void awaitThreadWaiting(Thread thread) throws InterruptedException { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TIMEOUT_SECONDS); + while (thread.getState() != Thread.State.WAITING && System.nanoTime() < deadline) { + Thread.sleep(1); + } + assertEquals("thread did not enter a wait", Thread.State.WAITING, thread.getState()); + } + + private static void awaitInterruptConsumed(Thread thread) throws InterruptedException { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TIMEOUT_SECONDS); + while (thread.isInterrupted() && System.nanoTime() < deadline) { + Thread.sleep(1); + } + assertTrue("caller returned before its accepted work", thread.isAlive()); + assertFalse("caller did not consume its interrupt", thread.isInterrupted()); + } + + private static void awaitUninterruptibly(CountDownLatch latch) { + boolean interrupted = false; + while (true) { + try { + latch.await(); + break; + } catch (InterruptedException ignored) { + interrupted = true; + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + + private static void shutdown(ExecutorService executor) throws InterruptedException { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } + + private static final class BlockingTask implements Callable { + private final CountDownLatch started = new CountDownLatch(1); + private final CountDownLatch release = new CountDownLatch(1); + + @Override + public Void call() { + started.countDown(); + awaitUninterruptibly(release); + return null; + } + + private void awaitStarted() throws InterruptedException { + assertTrue(started.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } + + private void awaitStartedUninterruptibly() { + awaitUninterruptibly(started); + } + + private void release() { + release.countDown(); + } + } + + @FunctionalInterface + private interface ThrowingAction { + void run() throws Throwable; + } + + private record ConcurrentRun(AtomicIntegerArray calls, int maxActiveHelpers) {} + + private record WorkerObservation( + Thread thread, String inheritedState, ClassLoader contextClassLoader) {} +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java new file mode 100644 index 0000000000..4a6034b210 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java @@ -0,0 +1,153 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; +import org.apache.lucene.util.hnsw.NeighborArray; +import org.junit.Test; + +/** Verifies serial and parallel CAGRA-adjacency materialization and host-copy ownership. */ +public class TestParallelGraphMaterialization extends LuceneTestCase { + + private static final int NUM_NODES = GPUBuiltHnswGraph.PARALLEL_MIN_NODES + 1000; + private static final int DEGREE = 12; + private static final int GRAPH_THREADS = 4; + + @Test + public void parallelMaterializationMatchesSerial() throws Exception { + int[][] rows = IntGraphTestMatrix.randomRows(NUM_NODES, DEGREE, 1); + IntGraphTestMatrix.ParallelExecutionProbe executionProbe = + new IntGraphTestMatrix.ParallelExecutionProbe(); + try (CuVSMatrix serialMatrix = new IntGraphTestMatrix(rows); + CuVSMatrix parallelMatrix = + new IntGraphTestMatrix.TrackingHostMatrix( + rows, new AtomicInteger(), null, executionProbe)) { + GPUBuiltHnswGraph serial = + new GPUBuiltHnswGraph( + NUM_NODES, /* dimensions= */ 4, Arrays.asList((int[]) null), List.of(serialMatrix)); + GPUBuiltHnswGraph parallel = newSingleLayerGraph(parallelMatrix, GRAPH_THREADS); + assertGraphsEqual(serial, parallel); + assertTrue(executionProbe.threadCount() > 1); + } + } + + @Test + public void overflowingDeviceShapeUsesSerialFallback() throws Exception { + int[][] adjacency = IntGraphTestMatrix.randomRows(NUM_NODES, 1, 0); + try (CuVSMatrix matrix = new IntGraphTestMatrix.DeviceMatrix(adjacency, Long.MAX_VALUE)) { + GPUBuiltHnswGraph graph = newSingleLayerGraph(matrix, GRAPH_THREADS); + assertEquals(NUM_NODES, graph.size()); + } + } + + @Test + public void failedDeviceCopyClosesHostAllocationAndSuppressesCloseFailure() { + RuntimeException copyFailure = new RuntimeException("copy failed"); + RuntimeException closeFailure = new RuntimeException("close failed"); + AtomicInteger hostCloseCount = new AtomicInteger(); + CuVSDeviceMatrix source = + new IntGraphTestMatrix.DeviceMatrix(new int[][] {{0}}, 1) { + @Override + public void toHost(CuVSHostMatrix target) { + throw copyFailure; + } + }; + CuVSHostMatrix hostCopy = + new IntGraphTestMatrix.TrackingHostMatrix(hostCloseCount, closeFailure); + + RuntimeException thrown = + assertThrows( + RuntimeException.class, () -> GPUBuiltHnswGraph.copyToHost(source, () -> hostCopy)); + + assertSame(copyFailure, thrown); + assertEquals(1, hostCloseCount.get()); + assertArrayEquals(new Throwable[] {closeFailure}, thrown.getSuppressed()); + } + + @Test + public void admittedDeviceCopyIsMaterializedInParallelAndReleased() throws Exception { + int[][] expectedRows = IntGraphTestMatrix.randomRows(NUM_NODES, DEGREE, 3); + int[][] wrongSourceRows = new int[NUM_NODES][DEGREE]; + AtomicInteger hostCloseCount = new AtomicInteger(); + AtomicInteger copyCount = new AtomicInteger(); + IntGraphTestMatrix.ParallelExecutionProbe executionProbe = + new IntGraphTestMatrix.ParallelExecutionProbe(); + CuVSHostMatrix hostCopy = + new IntGraphTestMatrix.TrackingHostMatrix( + expectedRows, hostCloseCount, null, executionProbe); + CuVSDeviceMatrix source = + new IntGraphTestMatrix.DeviceMatrix(wrongSourceRows, DEGREE) { + @Override + public void toHost(CuVSHostMatrix target) { + assertSame(hostCopy, target); + copyCount.incrementAndGet(); + } + }; + long requiredHeadroom = GraphCopyMemoryBudget.requiredHeadroom(NUM_NODES, DEGREE); + GraphCopyMemoryBudget budget = + new GraphCopyMemoryBudget( + () -> new GraphCopyMemoryBudget.MemorySnapshot(requiredHeadroom, requiredHeadroom)); + + NeighborArray[] neighbors = + GPUBuiltHnswGraph.materializeDeviceAdjacency( + source, NUM_NODES, GRAPH_THREADS, budget, () -> hostCopy); + + for (int node = 0; node < NUM_NODES; node++) { + assertArrayEquals( + expectedRows[node], Arrays.copyOf(neighbors[node].nodes(), neighbors[node].size())); + } + assertEquals(1, copyCount.get()); + assertEquals(1, hostCloseCount.get()); + assertTrue(executionProbe.threadCount() > 1); + try (GraphCopyMemoryBudget.Reservation ignored = + budget.tryReserve(NUM_NODES, DEGREE).orElseThrow()) { + // The first reservation was released after materialization. + } + } + + private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, int numThreads) + throws IOException { + return new GPUBuiltHnswGraph( + NUM_NODES, + /* dimensions= */ 4, + Arrays.asList((int[]) null), + List.of(layer0Adjacency), + numThreads); + } + + private static void assertGraphsEqual(HnswGraph a, HnswGraph b) throws Exception { + assertEquals(a.numLevels(), b.numLevels()); + for (int level = 0; level < a.numLevels(); level++) { + int[] nodes = NodesIterator.getSortedNodes(a.getNodesOnLevel(level)); + for (int node : nodes) { + assertArrayEquals( + "node " + node + " at level " + level + " has different neighbors", + arcsOf(a, level, node), + arcsOf(b, level, node)); + } + } + } + + private static int[] arcsOf(HnswGraph graph, int level, int node) throws Exception { + graph.seek(level, node); + List arcs = new ArrayList<>(); + for (int n = graph.nextNeighbor(); n != NO_MORE_DOCS; n = graph.nextNeighbor()) { + arcs.add(n); + } + return arcs.stream().mapToInt(Integer::intValue).toArray(); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java new file mode 100644 index 0000000000..e86d3deb36 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java @@ -0,0 +1,192 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CuVSMatrix; +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.hnsw.NeighborArray; +import org.junit.Test; + +/** Verifies parallel level-zero graph serialization is byte-identical to serial serialization. */ +public class TestParallelGraphSerialization extends LuceneTestCase { + + private static final int NUM_NODES = GPUBuiltHnswGraph.PARALLEL_MIN_NODES + 1000; + private static final int DEGREE = 12; + private static final int GRAPH_THREADS = 4; + + @Test + public void parallelSerializationMatchesSerial() throws Exception { + try (CuVSMatrix matrix = IntGraphTestMatrix.random(NUM_NODES, DEGREE, 2); + Directory dir = new ByteBuffersDirectory()) { + GPUBuiltHnswGraph serialGraph = newSingleLayerGraph(matrix); + IntGraphTestMatrix.ParallelExecutionProbe executionProbe = + new IntGraphTestMatrix.ParallelExecutionProbe(); + GPUBuiltHnswGraph parallelGraph = new RecordingGraph(matrix, executionProbe); + assertSerialAndParallelMatch(serialGraph, parallelGraph, dir); + assertTrue(executionProbe.threadCount() > 1); + } + } + + @Test + public void parallelSerializationMatchesSerialAcrossByteBoundedWave() throws Exception { + int maxConn = 512; + int waveNodes = AcceleratedHNSWUtils.serializationWaveNodes(maxConn); + int numNodes = Math.max(GPUBuiltHnswGraph.PARALLEL_MIN_NODES + 1, waveNodes + 2); + long worstCaseNodeBytes = 5L + maxConn * 5L; + IntGraphTestMatrix.ParallelExecutionProbe executionProbe = + new IntGraphTestMatrix.ParallelExecutionProbe(); + + try (Directory dir = new ByteBuffersDirectory()) { + assertTrue(waveNodes < AcceleratedHNSWUtils.MAX_SERIALIZATION_WAVE_NODES); + assertTrue( + waveNodes * worstCaseNodeBytes <= AcceleratedHNSWUtils.MAX_SERIALIZED_BYTES_PER_WAVE); + assertTrue( + (waveNodes + 1L) * worstCaseNodeBytes + > AcceleratedHNSWUtils.MAX_SERIALIZED_BYTES_PER_WAVE); + assertSerialAndParallelMatch( + new LazyBoundaryGraph(numNodes, maxConn, waveNodes, null), + new LazyBoundaryGraph(numNodes, maxConn, waveNodes, executionProbe), + dir); + assertTrue(executionProbe.threadCount() > 1); + } + } + + private static void assertSerialAndParallelMatch( + GPUBuiltHnswGraph serialGraph, GPUBuiltHnswGraph parallelGraph, Directory dir) + throws Exception { + int[][] serialOffsets; + try (IndexOutput out = dir.createOutput("serial", IOContext.DEFAULT)) { + serialOffsets = AcceleratedHNSWUtils.writeGraph(serialGraph, out); + } + int[][] parallelOffsets; + try (IndexOutput out = dir.createOutput("parallel", IOContext.DEFAULT)) { + parallelOffsets = AcceleratedHNSWUtils.writeGraph(parallelGraph, out, GRAPH_THREADS); + } + + assertEquals(serialOffsets.length, parallelOffsets.length); + for (int level = 0; level < serialOffsets.length; level++) { + assertArrayEquals(serialOffsets[level], parallelOffsets[level]); + } + assertArrayEquals(readAllBytes(dir, "serial"), readAllBytes(dir, "parallel")); + } + + private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency) { + return new GPUBuiltHnswGraph( + NUM_NODES, /* dimensions= */ 4, Arrays.asList((int[]) null), List.of(layer0Adjacency)); + } + + private static byte[] readAllBytes(Directory dir, String name) throws Exception { + try (IndexInput in = dir.openInput(name, IOContext.DEFAULT)) { + byte[] bytes = new byte[(int) in.length()]; + in.readBytes(bytes, 0, bytes.length); + return bytes; + } + } + + private static final class RecordingGraph extends GPUBuiltHnswGraph { + private final IntGraphTestMatrix.ParallelExecutionProbe executionProbe; + + RecordingGraph( + CuVSMatrix layer0Adjacency, IntGraphTestMatrix.ParallelExecutionProbe executionProbe) { + super(NUM_NODES, /* dimensions= */ 4, Arrays.asList((int[]) null), List.of(layer0Adjacency)); + this.executionProbe = executionProbe; + } + + @Override + public NeighborArray getNeighbors(int level, int node) { + executionProbe.recordExecution(); + return super.getNeighbors(level, node); + } + } + + /** Lazily makes boundary nodes distinct without retaining a heap graph. */ + private static final class LazyBoundaryGraph extends GPUBuiltHnswGraph { + private final int graphSize; + private final int maxConn; + private final int waveNodes; + private final IntGraphTestMatrix.ParallelExecutionProbe executionProbe; + + LazyBoundaryGraph( + int graphSize, + int maxConn, + int waveNodes, + IntGraphTestMatrix.ParallelExecutionProbe executionProbe) + throws IOException { + super( + 0, + /* dimensions= */ 4, + Arrays.asList((int[]) null), + List.of(new IntGraphTestMatrix(new int[0][])), + 1); + this.graphSize = graphSize; + this.maxConn = maxConn; + this.waveNodes = waveNodes; + this.executionProbe = executionProbe; + } + + @Override + public int size() { + return graphSize; + } + + @Override + public int maxConn() { + return maxConn; + } + + @Override + public NodesIterator getNodesOnLevel(int level) { + return new RangeNodesIterator(level == 0 ? graphSize : 0); + } + + @Override + public NeighborArray getNeighbors(int level, int node) { + if (executionProbe != null) { + executionProbe.recordExecution(); + } + if (node < waveNodes - 1) { + return null; + } + NeighborArray neighbors = new NeighborArray(1, true); + neighbors.addInOrder(node, 1.0f); + return neighbors; + } + } + + private static final class RangeNodesIterator extends GPUBuiltHnswGraph.NodesIterator { + private int current = -1; + + RangeNodesIterator(int size) { + super(size); + } + + @Override + public boolean hasNext() { + return current + 1 < size; + } + + @Override + public int nextInt() { + return ++current; + } + + @Override + public int consume(int[] dest) { + int count = Math.min(dest.length, size - (current + 1)); + for (int i = 0; i < count; i++) { + dest[i] = ++current; + } + return count; + } + } +} From 03d28e150d765c66bb7a395dc64482ba93d438ff Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Wed, 30 Sep 2026 15:25:46 +0000 Subject: [PATCH 15/21] Verify parallel HNSW graph processing paths --- ci/test_lucene.sh | 3 + ci/test_lucene_prebuilt.sh | 3 + ...nvidia-cuvs-lucene-acceleratedhnswutils.md | 14 +- ...om-nvidia-cuvs-lucene-gpubuilthnswgraph.md | 20 +- ...ne-lucene99acceleratedhnswvectorswriter.md | 14 +- ...leratedhnswbinaryquantizedvectorswriter.md | 14 +- ...leratedhnswscalarquantizedvectorswriter.md | 14 +- .../cuvs/lucene/AcceleratedHNSWUtils.java | 81 ++++- .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 89 ++++- .../cuvs/lucene/GraphProcessingTrace.java | 71 ++++ .../Lucene99AcceleratedHNSWVectorsWriter.java | 17 +- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 17 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 17 +- .../TestGraphThreadsPersistedIndex.java | 225 +++++++++---- .../TestParallelGraphSerialization.java | 71 +++- ...TestWriterThreadsGraphMaterialization.java | 307 ------------------ .../TestWriterThreadsGraphSerialization.java | 261 --------------- .../TestWriterThreadsPersistedIndex.java | 126 ------- 18 files changed, 546 insertions(+), 818 deletions(-) create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphProcessingTrace.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java delete mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsPersistedIndex.java diff --git a/ci/test_lucene.sh b/ci/test_lucene.sh index 5478a834ae..a07488f9b1 100755 --- a/ci/test_lucene.sh +++ b/ci/test_lucene.sh @@ -11,6 +11,9 @@ set +e rapids-logger "Check GPU usage" nvidia-smi +# This GPU lane must execute the real cuVS writer-path sentinel rather than skip it. +export CUVS_TESTS_REQUIRE_GPU=1 + rapids-logger "Run cuvs-lucene build and tests" RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}" diff --git a/ci/test_lucene_prebuilt.sh b/ci/test_lucene_prebuilt.sh index 477ff75ad3..05f384b910 100755 --- a/ci/test_lucene_prebuilt.sh +++ b/ci/test_lucene_prebuilt.sh @@ -33,6 +33,9 @@ fi rapids-logger "Check GPU usage" nvidia-smi +# This GPU lane must execute the real cuVS writer-path sentinel rather than skip it. +export CUVS_TESTS_REQUIRE_GPU=1 + rapids-logger "Configuring conda strict channel priority" conda config --set channel_priority strict diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md index f9321443e4..47a7161773 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md @@ -44,7 +44,7 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( int dimensions, CuVSMatrix a Creates a multi-layer HNSW graph from a native matrix without copying the complete dataset to the Java heap. The list view copies only rows selected for an upper layer. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:215`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:239`_ ### writeGraph @@ -71,7 +71,7 @@ a 2D array of offsets | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:370`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:416`_ ### writeMeta @@ -100,7 +100,7 @@ Writes the meta information for the index. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:516`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:589`_ ### printInfoStream @@ -116,7 +116,7 @@ A utility method to print info/debugging messages using InfoStream. | --- | --- | | `msg` | the debugging message to print | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:598`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:671`_ ### writeEmpty @@ -138,7 +138,7 @@ Writes an empty meta information for the field. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:610`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:683`_ ### quantizeFloatVectorsToBinary @@ -161,7 +161,7 @@ Bits are packed: 8 dimensions per byte. A list of byte binary representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:623`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:696`_ ### quantizeFloatVectorsToScalar @@ -181,6 +181,6 @@ Scalar quantization. A list of byte scalar representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:665`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:738`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:35`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md index 239e1a3f33..eb0a751516 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md @@ -41,7 +41,7 @@ public NodesIterator getNodesOnLevel(int level) Get all nodes on a given level as node 0th ordinals. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:228`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:295`_ ### getNeighbors @@ -62,7 +62,7 @@ Get the neighbors for the node and the level it resides. an instance of NeighborArray -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:246`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:313`_ ### seek @@ -72,7 +72,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Move the pointer to exactly the given level's target. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:271`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:338`_ ### nextNeighbor @@ -82,7 +82,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Iterates over the neighbor list. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:281`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:348`_ ### entryNode @@ -92,7 +92,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns graph's entry point on the top level. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:312`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:379`_ ### maxConn @@ -102,7 +102,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap returns M, the maximum number of connections for a node. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:331`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:398`_ ### neighborCount @@ -112,7 +112,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns the neighbor count. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:346`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:413`_ ### size @@ -122,7 +122,7 @@ public int size() Returns the number of nodes in the graph. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:421`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:488`_ ### numLevels @@ -136,7 +136,7 @@ Returns the number of levels in the HNSW graph. the number of levels -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:430`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:497`_ ### dimensions @@ -150,6 +150,6 @@ Gets the vector dimension. the vector dimension -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:439`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:506`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:27`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md index 18a59572c5..ea991d5f91 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md @@ -37,7 +37,7 @@ Initializes `Lucene99AcceleratedHNSWVectorsWriter` | --- | --- | | `IOException` | IOException | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:89`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:90`_ ### addField @@ -47,7 +47,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:131`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:133`_ ### flush @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:223`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:230`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:366`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:377`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:375`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:386`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:395`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:406`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:405`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:416`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:55`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md index b355a3784f..2d058c70f6 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md @@ -37,7 +37,7 @@ Initializes `LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter` | --- | --- | | `IOException` | IOException | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:79`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:80`_ ### addField @@ -47,7 +47,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:127`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:129`_ ### flush @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:222`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:229`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:308`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:319`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:339`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:350`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:359`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:370`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:368`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:379`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:57`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md index 6ba5228a80..e4fd6e79ff 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md @@ -37,7 +37,7 @@ Initializes `LuceneAcceleratedHNSWScalarQuantizedVectorsWriter` | --- | --- | | `IOException` | IOException | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:89`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:90`_ ### addField @@ -47,7 +47,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:136`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:138`_ ### flush @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:249`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:256`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:333`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:344`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:364`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:375`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:384`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:395`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:393`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:404`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:56`_ diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 620e4464f6..5257af0509 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -105,6 +105,29 @@ private static GPUBuiltHnswGraph createMultiLayerHnswGraph( QuantizationType quantization, int graphThreads) throws Throwable { + return createMultiLayerHnswGraph( + size, + dimensions, + adjacencyListMatrix, + vectors, + hnswLayers, + params, + quantization, + graphThreads, + GraphProcessingTrace.disabled()); + } + + private static GPUBuiltHnswGraph createMultiLayerHnswGraph( + int size, + int dimensions, + CuVSMatrix adjacencyListMatrix, + List vectors, + int hnswLayers, + CagraIndexParams params, + QuantizationType quantization, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) + throws Throwable { int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); @@ -192,7 +215,8 @@ private static GPUBuiltHnswGraph createMultiLayerHnswGraph( } // The graph eagerly copies all adjacency rows, so generated upper matrices can now close. - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, graphThreads); + return new GPUBuiltHnswGraph( + size, dimensions, layerNodes, layerAdjacencies, graphThreads, graphProcessingTrace); } catch (Throwable t) { failure = t; throw t; @@ -240,6 +264,27 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( QuantizationType quantization, int graphThreads) throws Throwable { + return createMultiLayerHnswGraph( + dimensions, + adjacencyListMatrix, + vectorDataset, + hnswLayers, + params, + quantization, + graphThreads, + GraphProcessingTrace.disabled()); + } + + static GPUBuiltHnswGraph createMultiLayerHnswGraph( + int dimensions, + CuVSMatrix adjacencyListMatrix, + CuVSMatrix vectorDataset, + int hnswLayers, + CagraIndexParams params, + QuantizationType quantization, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) + throws Throwable { int size = Math.toIntExact(vectorDataset.size()); // Matrix columns are the stored width: binary vectors are bit-packed, while scalar and float // vectors store one value per dimension. @@ -272,7 +317,8 @@ public int size() { hnswLayers, params, quantization, - graphThreads); + graphThreads, + graphProcessingTrace); } private static Throwable closeUpperLayerAdjacencies(List layerAdjacencies) { @@ -375,6 +421,15 @@ public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorInde static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int graphThreads) throws IOException { + return writeGraph(graph, vectorIndex, graphThreads, GraphProcessingTrace.disabled()); + } + + static int[][] writeGraph( + GPUBuiltHnswGraph graph, + IndexOutput vectorIndex, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) + throws IOException { int countOnLevel0 = graph.size(); int numLevels = graph.numLevels(); int[][] offsets = new int[numLevels][]; @@ -385,8 +440,22 @@ static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int if (graphThreads > 1 && level0Nodes.length >= GPUBuiltHnswGraph.PARALLEL_MIN_NODES) { writeLevel0Parallel( graph, vectorIndex, level0Nodes, offsets[0], countOnLevel0, maxConn, graphThreads); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.SERIALIZATION, + GraphProcessingTrace.Mode.PARALLEL, + GraphProcessingTrace.Reason.ABOVE_THRESHOLD, + graphThreads, + level0Nodes.length); } else { writeLevelSerial(graph, vectorIndex, 0, level0Nodes, offsets[0], countOnLevel0, maxConn); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.SERIALIZATION, + GraphProcessingTrace.Mode.SERIAL, + graphThreads <= 1 + ? GraphProcessingTrace.Reason.SINGLE_THREAD + : GraphProcessingTrace.Reason.BELOW_THRESHOLD, + graphThreads, + level0Nodes.length); } for (int level = 1; level < numLevels; level++) { @@ -398,7 +467,11 @@ static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int return offsets; } - /** Absolute node cap for one serialization wave. */ + /** + * Fixed operational guardrail on nodes processed before task-local buffers are concatenated when + * the encoded-byte limit would otherwise permit a very large wave. Unlike the byte limit below, + * this is a policy cap rather than an encoded-size calculation. + */ static final int MAX_SERIALIZATION_WAVE_NODES = 1 << 20; /** Maximum worst-case encoded payload buffered by one serialization wave. */ @@ -481,7 +554,7 @@ static int serializationWaveNodes(int maxConn) { private static void encodeNode( NeighborArray neighbors, int[] scratch, DataOutput out, int countOnLevel0) throws IOException { - int size = neighbors == null ? 0 : neighbors.size(); + int size = neighbors.size(); int actualSize = 0; if (size > 0) { int[] nodes = neighbors.nodes(); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index 2354e11be8..61c73b10ab 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -67,7 +67,25 @@ public GPUBuiltHnswGraph( List layerAdjacencies, int graphThreads) throws IOException { - this(size, dimensions, materialize(size, layerNodes, layerAdjacencies, graphThreads)); + this( + size, + dimensions, + materialize( + size, layerNodes, layerAdjacencies, graphThreads, GraphProcessingTrace.disabled())); + } + + GPUBuiltHnswGraph( + int size, + int dimensions, + List layerNodes, + List layerAdjacencies, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) + throws IOException { + this( + size, + dimensions, + materialize(size, layerNodes, layerAdjacencies, graphThreads, graphProcessingTrace)); } private GPUBuiltHnswGraph(int size, int dimensions, MaterializedGraph graph) { @@ -82,7 +100,7 @@ private GPUBuiltHnswGraph(int size, int dimensions, MaterializedGraph graph) { private static MaterializedGraph materializeSerial( int size, List layerNodes, List layerAdjacencies) { try { - return materialize(size, layerNodes, layerAdjacencies, 1); + return materialize(size, layerNodes, layerAdjacencies, 1, GraphProcessingTrace.disabled()); } catch (IOException impossible) { throw new AssertionError( "serial graph materialization cannot fail with IOException", impossible); @@ -90,18 +108,23 @@ private static MaterializedGraph materializeSerial( } private static MaterializedGraph materialize( - int size, List layerNodes, List layerAdjacencies, int graphThreads) + int size, + List layerNodes, + List layerAdjacencies, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) throws IOException { List upperLayerNodes = new ArrayList<>(); List upperLayerNeighbors = new ArrayList<>(); NeighborArray[] baseLayerNeighbors = - fillNeighborArray(layerAdjacencies.get(0), size, graphThreads); + fillNeighborArray(layerAdjacencies.get(0), size, graphThreads, graphProcessingTrace); for (int level = 1; level < layerAdjacencies.size(); level++) { int[] nodes = layerNodes.get(level); upperLayerNodes.add(nodes); upperLayerNeighbors.add( - fillNeighborArray(layerAdjacencies.get(level), nodes.length, graphThreads)); + fillNeighborArray( + layerAdjacencies.get(level), nodes.length, graphThreads, graphProcessingTrace)); } return new MaterializedGraph( layerAdjacencies.size(), upperLayerNodes, baseLayerNeighbors, upperLayerNeighbors); @@ -116,10 +139,20 @@ private static MaterializedGraph materialize( * @param graphThreads maximum threads for this operation, including the caller * @return the NeighborArray */ - private static NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size, int graphThreads) + private static NeighborArray[] fillNeighborArray( + CuVSMatrix adjacency, int size, int graphThreads, GraphProcessingTrace graphProcessingTrace) throws IOException { if (graphThreads <= 1 || size < PARALLEL_MIN_NODES) { - return fillNeighborArraySerial(adjacency, size); + NeighborArray[] neighbors = fillNeighborArraySerial(adjacency, size); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.MATERIALIZATION, + GraphProcessingTrace.Mode.SERIAL, + graphThreads <= 1 + ? GraphProcessingTrace.Reason.SINGLE_THREAD + : GraphProcessingTrace.Reason.BELOW_THRESHOLD, + graphThreads, + size); + return neighbors; } if (adjacency instanceof CuVSDeviceMatrix deviceAdjacency) { @@ -128,10 +161,18 @@ private static NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size, size, graphThreads, GraphCopyMemoryBudget.system(), - () -> newHostMatrix(deviceAdjacency)); + () -> newHostMatrix(deviceAdjacency), + graphProcessingTrace); } - return fillNeighborArrayParallel(adjacency, size, graphThreads); + NeighborArray[] neighbors = fillNeighborArrayParallel(adjacency, size, graphThreads); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.MATERIALIZATION, + GraphProcessingTrace.Mode.PARALLEL, + GraphProcessingTrace.Reason.HOST_SOURCE, + graphThreads, + size); + return neighbors; } private static NeighborArray[] fillNeighborArraySerial(CuVSMatrix source, int size) { @@ -147,14 +188,40 @@ static NeighborArray[] materializeDeviceAdjacency( GraphCopyMemoryBudget memoryBudget, Supplier hostCopyFactory) throws IOException { + return materializeDeviceAdjacency( + source, size, graphThreads, memoryBudget, hostCopyFactory, GraphProcessingTrace.disabled()); + } + + private static NeighborArray[] materializeDeviceAdjacency( + CuVSDeviceMatrix source, + int size, + int graphThreads, + GraphCopyMemoryBudget memoryBudget, + Supplier hostCopyFactory, + GraphProcessingTrace graphProcessingTrace) + throws IOException { Optional reservation = memoryBudget.tryReserve(source.size(), source.columns()); if (reservation.isEmpty()) { - return fillNeighborArraySerial(source, size); + NeighborArray[] neighbors = fillNeighborArraySerial(source, size); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.MATERIALIZATION, + GraphProcessingTrace.Mode.SERIAL, + GraphProcessingTrace.Reason.MEMORY_ADMISSION_DENIED, + graphThreads, + size); + return neighbors; } try (GraphCopyMemoryBudget.Reservation ignored = reservation.orElseThrow(); CuVSHostMatrix hostCopy = copyToHost(source, hostCopyFactory)) { - return fillNeighborArrayParallel(hostCopy, size, graphThreads); + NeighborArray[] neighbors = fillNeighborArrayParallel(hostCopy, size, graphThreads); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.MATERIALIZATION, + GraphProcessingTrace.Mode.PARALLEL, + GraphProcessingTrace.Reason.DEVICE_HOST_COPY, + graphThreads, + size); + return neighbors; } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphProcessingTrace.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphProcessingTrace.java new file mode 100644 index 0000000000..54f271fd07 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphProcessingTrace.java @@ -0,0 +1,71 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.util.Locale; +import org.apache.lucene.util.InfoStream; + +/** Reports the graph-processing path selected for one Lucene writer. */ +final class GraphProcessingTrace { + + enum Stage { + MATERIALIZATION, + SERIALIZATION + } + + enum Mode { + SERIAL, + PARALLEL + } + + enum Reason { + ABOVE_THRESHOLD, + BELOW_THRESHOLD, + DEVICE_HOST_COPY, + HOST_SOURCE, + MEMORY_ADMISSION_DENIED, + SINGLE_THREAD + } + + private static final GraphProcessingTrace DISABLED = new GraphProcessingTrace(null, null); + + private final InfoStream infoStream; + private final String component; + + private GraphProcessingTrace(InfoStream infoStream, String component) { + this.infoStream = infoStream; + this.component = component; + } + + static GraphProcessingTrace disabled() { + return DISABLED; + } + + static GraphProcessingTrace toInfoStream(InfoStream infoStream, String component) { + return new GraphProcessingTrace(infoStream, component); + } + + void record(Stage stage, Mode mode, Reason reason, int requestedThreads, int nodes) { + if (infoStream == null || !infoStream.isEnabled(component)) { + return; + } + infoStream.message( + component, + "graph-processing stage=" + + lowerCase(stage) + + " mode=" + + lowerCase(mode) + + " reason=" + + lowerCase(reason) + + " requestedThreads=" + + requestedThreads + + " nodes=" + + nodes); + } + + private static String lowerCase(Enum value) { + return value.name().toLowerCase(Locale.ROOT).replace('_', '-'); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index f113f5e198..8f52eae6e0 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -64,6 +64,7 @@ public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; + private final GraphProcessingTrace graphProcessingTrace; private IndexOutput hnswMeta = null; private IndexOutput hnswVectorIndex = null; private String vemFileName; @@ -95,6 +96,7 @@ public Lucene99AcceleratedHNSWVectorsWriter( super(); this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; + this.graphProcessingTrace = GraphProcessingTrace.toInfoStream(infoStream, COMPONENT); this.acceleratedHNSWParams = acceleratedHNSWParams; vemFileName = IndexFileNames.segmentFileName( @@ -186,10 +188,15 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.NONE, - acceleratedHNSWParams.getGraphThreads()); + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getGraphThreads()); + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -275,7 +282,11 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getGraphThreads()); + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 3cf50be3b7..4fa02eb1d5 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -63,6 +63,7 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter extends KnnVector private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; + private final GraphProcessingTrace graphProcessingTrace; private final AcceleratedHNSWParams acceleratedHNSWParams; private IndexOutput hnswMeta = null, hnswVectorIndex = null; private boolean finished; @@ -86,6 +87,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter( this.acceleratedHNSWParams = acceleratedHNSWParams; this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; + this.graphProcessingTrace = GraphProcessingTrace.toInfoStream(infoStream, COMPONENT); vemFileName = IndexFileNames.segmentFileName( @@ -184,11 +186,16 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.BINARY, - acceleratedHNSWParams.getGraphThreads()); + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getGraphThreads()); + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -283,7 +290,11 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { // Write the graph to the vector index int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getGraphThreads()); + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 77370fb26b..0ad8f24c5e 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -64,6 +64,7 @@ public class LuceneAcceleratedHNSWScalarQuantizedVectorsWriter extends KnnVector private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; + private final GraphProcessingTrace graphProcessingTrace; private final AcceleratedHNSWParams acceleratedHNSWParams; private IndexOutput hnswMeta = null, hnswVectorIndex = null; private boolean finished; @@ -96,6 +97,7 @@ public LuceneAcceleratedHNSWScalarQuantizedVectorsWriter( this.acceleratedHNSWParams = acceleratedHNSWParams; this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; + this.graphProcessingTrace = GraphProcessingTrace.toInfoStream(infoStream, COMPONENT); vemFileName = IndexFileNames.segmentFileName( @@ -211,11 +213,16 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.SCALAR, - acceleratedHNSWParams.getGraphThreads()); + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getGraphThreads()); + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -309,7 +316,11 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getGraphThreads()); + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphThreadsPersistedIndex.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphThreadsPersistedIndex.java index cf980d439b..e951accc7e 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphThreadsPersistedIndex.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphThreadsPersistedIndex.java @@ -8,6 +8,11 @@ import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import com.carrotsearch.randomizedtesting.annotations.Name; +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import java.util.Random; import org.apache.lucene.codecs.Codec; import org.apache.lucene.codecs.KnnVectorsReader; @@ -29,6 +34,7 @@ import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.hnsw.HnswGraph; import org.junit.Test; @@ -36,87 +42,162 @@ public class TestGraphThreadsPersistedIndex extends LuceneTestCase { private static final String VECTOR_FIELD = "vector"; + private static final String REQUIRE_GPU_ENV = "CUVS_TESTS_REQUIRE_GPU"; private static final int VECTOR_COUNT = GPUBuiltHnswGraph.PARALLEL_MIN_NODES + 1; - private static final int DIMENSIONS = 32; + private static final int GRAPH_THREADS = 4; + + private enum WriterVariant { + FLOAT, + BINARY_QUANTIZED, + SCALAR_QUANTIZED + } + + private final WriterVariant writerVariant; + private final int dimensions; + + public TestGraphThreadsPersistedIndex( + @Name("writer") WriterVariant writerVariant, @Name("dimensions") int dimensions) { + this.writerVariant = writerVariant; + this.dimensions = dimensions; + } + + @ParametersFactory + public static List parameters() { + return List.of( + new Object[] {WriterVariant.FLOAT, 32}, + new Object[] {WriterVariant.BINARY_QUANTIZED, 129}, + new Object[] {WriterVariant.SCALAR_QUANTIZED, 32}); + } @Test - public void testGraphThreadsRoundTripAboveThreshold() throws Exception { - assumeTrue("cuVS not supported", isSupported()); + public void testEveryWriterForwardsGraphThreadsAndPersistsValidIndex() throws Exception { + requireCuvsSupport(); AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder() .withWriterThreads(1) - .withGraphThreads(4) + .withGraphThreads(GRAPH_THREADS) .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) .withIntermediateGraphDegree(32) .withGraphDegree(16) .withHNSWLayer(1) .build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params); - Random random = new Random(0x2594L); + Codec codec = codecFor(params); + RecordingInfoStream infoStream = new RecordingInfoStream(); try (Directory directory = newDirectory()) { - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(VECTOR_COUNT + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); - try (IndexWriter writer = new IndexWriter(directory, config)) { - for (int id = 0; id < VECTOR_COUNT; id++) { - float[] vector = new float[DIMENSIONS]; - for (int dimension = 0; dimension < DIMENSIONS; dimension++) { - vector[dimension] = random.nextFloat(); - } - Document document = new Document(); - document.add(new StringField("id", Integer.toString(id), Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, vector, EUCLIDEAN)); - writer.addDocument(document); + writeIndexWithGraphThreads(directory, codec, infoStream); + assertParallelStageObserved(infoStream, "materialization", "device-host-copy"); + assertParallelStageObserved(infoStream, "serialization", "above-threshold"); + assertPersistedGraphIsValidAndSearchable(directory); + } + } + + private void writeIndexWithGraphThreads( + Directory directory, Codec codec, RecordingInfoStream infoStream) throws Exception { + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setInfoStream(infoStream) + .setUseCompoundFile(false) + .setMaxBufferedDocs(VECTOR_COUNT + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + Random random = new Random(0x2594L); + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int id = 0; id < VECTOR_COUNT; id++) { + float[] vector = new float[dimensions]; + for (int dimension = 0; dimension < dimensions; dimension++) { + vector[dimension] = random.nextFloat(); } + Document document = new Document(); + document.add(new StringField("id", Integer.toString(id), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, vector, EUCLIDEAN)); + writer.addDocument(document); } + } + } - TestUtil.checkIndex(directory); - try (DirectoryReader reader = DirectoryReader.open(directory)) { - assertEquals(1, reader.leaves().size()); - assertEquals(VECTOR_COUNT, reader.numDocs()); - LeafReader leaf = getOnlyLeafReader(reader); - HnswGraph graph = graphOf(leaf); - assertEquals(VECTOR_COUNT, graph.size()); - int arcs = 0; - HnswGraph.NodesIterator nodes = graph.getNodesOnLevel(0); - while (nodes.hasNext()) { - int node = nodes.nextInt(); - graph.seek(0, node); - for (int neighbor = graph.nextNeighbor(); - neighbor != NO_MORE_DOCS; - neighbor = graph.nextNeighbor()) { - assertTrue(neighbor >= 0); - assertTrue(neighbor < VECTOR_COUNT); - arcs++; - } + private static void assertPersistedGraphIsValidAndSearchable(Directory directory) + throws Exception { + TestUtil.checkIndex(directory); + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals(1, reader.leaves().size()); + assertEquals(VECTOR_COUNT, reader.numDocs()); + LeafReader leaf = getOnlyLeafReader(reader); + HnswGraph graph = graphOf(leaf); + assertEquals(VECTOR_COUNT, graph.size()); + int arcs = 0; + HnswGraph.NodesIterator nodes = graph.getNodesOnLevel(0); + while (nodes.hasNext()) { + int node = nodes.nextInt(); + graph.seek(0, node); + for (int neighbor = graph.nextNeighbor(); + neighbor != NO_MORE_DOCS; + neighbor = graph.nextNeighbor()) { + assertTrue(neighbor >= 0); + assertTrue(neighbor < VECTOR_COUNT); + arcs++; } - assertTrue("persisted graph contains no arcs", arcs > 0); - - int queryNode = graph.entryNode(); - assertTrue(queryNode >= 0); - assertTrue(queryNode < VECTOR_COUNT); - FloatVectorValues values = leaf.getFloatVectorValues(VECTOR_FIELD); - assertNotNull(values); - float[] query = values.vectorValue(queryNode).clone(); - int queryDoc = values.ordToDoc(queryNode); - String queryId = leaf.storedFields().document(queryDoc).get("id"); - - IndexSearcher searcher = new IndexSearcher(reader); - var hits = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, query, 10), 10); - assertEquals(10, hits.scoreDocs.length); - boolean foundQueryNode = false; - for (var hit : hits.scoreDocs) { - foundQueryNode |= queryId.equals(searcher.storedFields().document(hit.doc).get("id")); - } - assertTrue("the entry-node vector must be returned for its own query", foundQueryNode); } + assertTrue("persisted graph contains no arcs", arcs > 0); + + int queryNode = graph.entryNode(); + assertTrue(queryNode >= 0); + assertTrue(queryNode < VECTOR_COUNT); + FloatVectorValues values = leaf.getFloatVectorValues(VECTOR_FIELD); + assertNotNull(values); + float[] query = values.vectorValue(queryNode).clone(); + int queryDoc = values.ordToDoc(queryNode); + String queryId = leaf.storedFields().document(queryDoc).get("id"); + + IndexSearcher searcher = new IndexSearcher(reader); + var hits = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, query, 10), 10); + assertEquals(10, hits.scoreDocs.length); + boolean foundQueryNode = false; + for (var hit : hits.scoreDocs) { + foundQueryNode |= queryId.equals(searcher.storedFields().document(hit.doc).get("id")); + } + assertTrue("the entry-node vector must be returned for its own query", foundQueryNode); } } + private static void requireCuvsSupport() { + boolean supported = isSupported(); + if ("1".equals(System.getenv(REQUIRE_GPU_ENV))) { + assertTrue( + REQUIRE_GPU_ENV + + "=1, but cuVS failed to initialize; verify GPU visibility, matching " + + "libcuvs/libcuvs_c libraries, LD_LIBRARY_PATH, and Java native-access " + + "configuration.", + supported); + } else { + assumeTrue("cuVS not supported", supported); + } + } + + private Codec codecFor(AcceleratedHNSWParams params) throws Exception { + return switch (writerVariant) { + case FLOAT -> new Lucene101AcceleratedHNSWCodec(params); + case BINARY_QUANTIZED -> new LuceneAcceleratedHNSWBinaryQuantizedCodec(params); + case SCALAR_QUANTIZED -> new LuceneAcceleratedHNSWScalarQuantizedCodec(params); + }; + } + + private void assertParallelStageObserved( + RecordingInfoStream infoStream, String stage, String expectedReason) { + String expectedPath = + "graph-processing stage=" + + stage + + " mode=parallel reason=" + + expectedReason + + " requestedThreads=" + + GRAPH_THREADS + + " nodes=" + + VECTOR_COUNT; + assertTrue( + writerVariant + " did not report the expected path; messages: " + infoStream.messages(), + infoStream.messages().stream().anyMatch(message -> message.contains(expectedPath))); + } + private static HnswGraph graphOf(LeafReader leaf) throws Exception { KnnVectorsReader reader = ((CodecReader) leaf).getVectorReader(); if (reader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { @@ -124,4 +205,28 @@ private static HnswGraph graphOf(LeafReader leaf) throws Exception { } return ((HnswGraphProvider) reader).getGraph(VECTOR_FIELD); } + + private static final class RecordingInfoStream extends InfoStream { + + private final List messages = Collections.synchronizedList(new ArrayList<>()); + + @Override + public void message(String component, String message) { + messages.add(component + ": " + message); + } + + @Override + public boolean isEnabled(String component) { + return true; + } + + @Override + public void close() {} + + List messages() { + synchronized (messages) { + return List.copyOf(messages); + } + } + } } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java index e86d3deb36..ac6c1fecc4 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java @@ -61,6 +61,35 @@ public void parallelSerializationMatchesSerialAcrossByteBoundedWave() throws Exc } } + @Test + public void lowDegreeSerializationWaveUsesAbsoluteNodeCap() { + assertEquals( + AcceleratedHNSWUtils.MAX_SERIALIZATION_WAVE_NODES, + AcceleratedHNSWUtils.serializationWaveNodes(/* maxConn= */ 0)); + } + + @Test + public void serialSerializationRejectsMissingAdjacency() throws Exception { + assertMissingAdjacencyRejected(/* graphSize= */ 1, /* graphThreads= */ 1); + } + + @Test + public void parallelSerializationRejectsMissingAdjacency() throws Exception { + assertMissingAdjacencyRejected( + GPUBuiltHnswGraph.PARALLEL_MIN_NODES, /* graphThreads= */ GRAPH_THREADS); + } + + private static void assertMissingAdjacencyRejected(int graphSize, int graphThreads) + throws Exception { + try (Directory dir = new ByteBuffersDirectory(); + IndexOutput out = dir.createOutput("missing-adjacency", IOContext.DEFAULT)) { + GPUBuiltHnswGraph graph = new MissingAdjacencyGraph(graphSize); + expectThrows( + NullPointerException.class, + () -> AcceleratedHNSWUtils.writeGraph(graph, out, graphThreads)); + } + } + private static void assertSerialAndParallelMatch( GPUBuiltHnswGraph serialGraph, GPUBuiltHnswGraph parallelGraph, Directory dir) throws Exception { @@ -109,8 +138,12 @@ public NeighborArray getNeighbors(int level, int node) { } } - /** Lazily makes boundary nodes distinct without retaining a heap graph. */ + /** + * Crosses a maxConn-derived byte-bounded wave with sparse rows and no retained heap graph. + */ private static final class LazyBoundaryGraph extends GPUBuiltHnswGraph { + private static final NeighborArray EMPTY_NEIGHBORS = new NeighborArray(0, true); + private final int graphSize; private final int maxConn; private final int waveNodes; @@ -155,7 +188,7 @@ public NeighborArray getNeighbors(int level, int node) { executionProbe.recordExecution(); } if (node < waveNodes - 1) { - return null; + return EMPTY_NEIGHBORS; } NeighborArray neighbors = new NeighborArray(1, true); neighbors.addInOrder(node, 1.0f); @@ -163,6 +196,40 @@ public NeighborArray getNeighbors(int level, int node) { } } + private static final class MissingAdjacencyGraph extends GPUBuiltHnswGraph { + private final int graphSize; + + MissingAdjacencyGraph(int graphSize) throws IOException { + super( + 0, + /* dimensions= */ 4, + Arrays.asList((int[]) null), + List.of(new IntGraphTestMatrix(new int[0][])), + 1); + this.graphSize = graphSize; + } + + @Override + public int size() { + return graphSize; + } + + @Override + public int maxConn() { + return 0; + } + + @Override + public NodesIterator getNodesOnLevel(int level) { + return new RangeNodesIterator(level == 0 ? graphSize : 0); + } + + @Override + public NeighborArray getNeighbors(int level, int node) { + return null; + } + } + private static final class RangeNodesIterator extends GPUBuiltHnswGraph.NodesIterator { private int current = -1; diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java deleted file mode 100644 index ad3366dd14..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java +++ /dev/null @@ -1,307 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; - -import com.nvidia.cuvs.CuVSDeviceMatrix; -import com.nvidia.cuvs.CuVSHostMatrix; -import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.RowView; -import java.io.IOException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; -import java.util.Random; -import java.util.concurrent.atomic.AtomicInteger; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.util.hnsw.HnswGraph; -import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; -import org.junit.Test; - -/** Verifies serial and parallel CAGRA-adjacency materialization are equivalent and bounded. */ -public class TestWriterThreadsGraphMaterialization extends LuceneTestCase { - - private static final int NUM_NODES = GPUBuiltHnswGraph.PARALLEL_MIN_NODES + 1000; - private static final int DEGREE = 12; - private static final int NUM_THREADS = 4; - - @Test - public void parallelMaterializationMatchesSerial() throws Exception { - int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(1)); - - try (CuVSMatrix matrix = new ArrayMatrix(adjacency)) { - GPUBuiltHnswGraph serial = - new GPUBuiltHnswGraph( - NUM_NODES, /* dimensions= */ 4, Arrays.asList((int[]) null), List.of(matrix)); - GPUBuiltHnswGraph parallel = newSingleLayerGraph(matrix, NUM_THREADS); - assertGraphsEqual(serial, parallel); - } - } - - @Test - public void graphCopyBudgetHandlesExpectedDatasetSizesAndOverflow() { - assertTrue(GPUBuiltHnswGraph.fitsParallelGraphCopyBudget(25_000_000L, 32)); - assertFalse(GPUBuiltHnswGraph.fitsParallelGraphCopyBudget(100_000_000L, 32)); - assertFalse(GPUBuiltHnswGraph.fitsParallelGraphCopyBudget(Long.MAX_VALUE, Long.MAX_VALUE)); - } - - @Test - public void oversizedDeviceAdjacencyUsesSerialFallback() throws Exception { - int[][] adjacency = randomAdjacency(NUM_NODES, 1, new Random(0)); - long oversizedColumns = - GPUBuiltHnswGraph.MAX_PARALLEL_GRAPH_COPY_BYTES / Integer.BYTES / NUM_NODES + 1; - try (CuVSMatrix matrix = new ArrayDeviceMatrix(adjacency, oversizedColumns)) { - GPUBuiltHnswGraph graph = newSingleLayerGraph(matrix, NUM_THREADS); - assertEquals(NUM_NODES, graph.size()); - } - } - - @Test - public void failedDeviceCopyClosesHostAllocationAndSuppressesCloseFailure() { - RuntimeException copyFailure = new RuntimeException("copy failed"); - RuntimeException closeFailure = new RuntimeException("close failed"); - AtomicInteger hostCloseCount = new AtomicInteger(); - CuVSDeviceMatrix source = - new ArrayDeviceMatrix(new int[][] {{0}}, 1) { - @Override - public void toHost(CuVSHostMatrix target) { - throw copyFailure; - } - }; - CuVSHostMatrix hostCopy = new TrackingHostMatrix(hostCloseCount, closeFailure); - - RuntimeException thrown = - assertThrows( - RuntimeException.class, () -> GPUBuiltHnswGraph.copyToHost(source, () -> hostCopy)); - - assertSame(copyFailure, thrown); - assertEquals(1, hostCloseCount.get()); - assertArrayEquals(new Throwable[] {closeFailure}, thrown.getSuppressed()); - } - - @Test - public void successfulDeviceCopyTransfersHostOwnershipToCaller() { - AtomicInteger hostCloseCount = new AtomicInteger(); - CuVSHostMatrix hostCopy = new TrackingHostMatrix(hostCloseCount, null); - CuVSDeviceMatrix source = - new ArrayDeviceMatrix(new int[][] {{0}}, 1) { - @Override - public void toHost(CuVSHostMatrix target) { - assertSame(hostCopy, target); - } - }; - - CuVSHostMatrix returned = GPUBuiltHnswGraph.copyToHost(source, () -> hostCopy); - - assertSame(hostCopy, returned); - assertEquals(0, hostCloseCount.get()); - returned.close(); - assertEquals(1, hostCloseCount.get()); - } - - private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, int numThreads) - throws IOException { - return new GPUBuiltHnswGraph( - NUM_NODES, - /* dimensions= */ 4, - Arrays.asList((int[]) null), - List.of(layer0Adjacency), - numThreads); - } - - private static void assertGraphsEqual(HnswGraph a, HnswGraph b) throws Exception { - assertEquals(a.numLevels(), b.numLevels()); - for (int level = 0; level < a.numLevels(); level++) { - int[] nodes = NodesIterator.getSortedNodes(a.getNodesOnLevel(level)); - for (int node : nodes) { - assertArrayEquals( - "node " + node + " at level " + level + " has different neighbors", - arcsOf(a, level, node), - arcsOf(b, level, node)); - } - } - } - - private static int[] arcsOf(HnswGraph graph, int level, int node) throws Exception { - graph.seek(level, node); - List arcs = new ArrayList<>(); - for (int n = graph.nextNeighbor(); n != NO_MORE_DOCS; n = graph.nextNeighbor()) { - arcs.add(n); - } - return arcs.stream().mapToInt(Integer::intValue).toArray(); - } - - private static int[][] randomAdjacency(int numNodes, int degree, Random random) { - int[][] adjacency = new int[numNodes][degree]; - for (int[] row : adjacency) { - for (int j = 0; j < degree; j++) { - row[j] = random.nextInt(numNodes); - } - } - return adjacency; - } - - private static class ArrayMatrix implements CuVSMatrix { - private final int[][] rows; - - ArrayMatrix(int[][] rows) { - this.rows = rows; - } - - @Override - public long size() { - return rows.length; - } - - @Override - public long columns() { - return rows.length == 0 ? 0 : rows[0].length; - } - - @Override - public DataType dataType() { - return DataType.INT; - } - - @Override - public RowView getRow(long row) { - return new ArrayRow(rows[Math.toIntExact(row)]); - } - - @Override - public void toArray(int[][] target) { - for (int i = 0; i < rows.length; i++) { - System.arraycopy(rows[i], 0, target[i], 0, rows[i].length); - } - } - - @Override - public void toArray(float[][] target) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(byte[][] target) { - throw new UnsupportedOperationException(); - } - - @Override - public void toHost(CuVSHostMatrix target) { - throw new UnsupportedOperationException(); - } - - @Override - public CuVSHostMatrix toHost() { - throw new UnsupportedOperationException(); - } - - @Override - public void toDevice(CuVSDeviceMatrix target, CuVSResources resources) { - throw new UnsupportedOperationException(); - } - - @Override - public CuVSDeviceMatrix toDevice(CuVSResources resources) { - throw new UnsupportedOperationException(); - } - - @Override - public void close() {} - } - - /** Reports an oversized device shape and fails if the guarded host-copy path is reached. */ - private static class ArrayDeviceMatrix extends ArrayMatrix implements CuVSDeviceMatrix { - private final long reportedColumns; - - ArrayDeviceMatrix(int[][] rows, long reportedColumns) { - super(rows); - this.reportedColumns = reportedColumns; - } - - @Override - public long columns() { - return reportedColumns; - } - - @Override - public void toHost(CuVSHostMatrix target) { - throw new AssertionError("oversized device adjacency must not be copied to host"); - } - - @Override - public CuVSHostMatrix toHost() { - throw new AssertionError("oversized device adjacency must not be copied to host"); - } - } - - private static final class TrackingHostMatrix extends ArrayMatrix implements CuVSHostMatrix { - private final AtomicInteger closeCount; - private final RuntimeException closeFailure; - - TrackingHostMatrix(AtomicInteger closeCount, RuntimeException closeFailure) { - super(new int[][] {{0}}); - this.closeCount = closeCount; - this.closeFailure = closeFailure; - } - - @Override - public int get(int row, int column) { - return 0; - } - - @Override - public void close() { - closeCount.incrementAndGet(); - if (closeFailure != null) { - throw closeFailure; - } - } - } - - private static final class ArrayRow implements RowView { - private final int[] values; - - ArrayRow(int[] values) { - this.values = values; - } - - @Override - public long size() { - return values.length; - } - - @Override - public int getAsInt(long index) { - return values[Math.toIntExact(index)]; - } - - @Override - public float getAsFloat(long index) { - throw new UnsupportedOperationException(); - } - - @Override - public byte getAsByte(long index) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(int[] target) { - System.arraycopy(values, 0, target, 0, values.length); - } - - @Override - public void toArray(float[] target) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(byte[] target) { - throw new UnsupportedOperationException(); - } - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java deleted file mode 100644 index 2d7188b3a1..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java +++ /dev/null @@ -1,261 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import com.nvidia.cuvs.CuVSDeviceMatrix; -import com.nvidia.cuvs.CuVSHostMatrix; -import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.RowView; -import java.io.IOException; -import java.util.Arrays; -import java.util.List; -import java.util.Random; -import org.apache.lucene.store.ByteBuffersDirectory; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.IOContext; -import org.apache.lucene.store.IndexInput; -import org.apache.lucene.store.IndexOutput; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.util.hnsw.NeighborArray; -import org.junit.Test; - -/** Verifies parallel level-zero graph serialization is byte-identical to serial serialization. */ -public class TestWriterThreadsGraphSerialization extends LuceneTestCase { - - private static final int NUM_NODES = AcceleratedHNSWUtils.PARALLEL_MIN_NODES + 1000; - private static final int DEGREE = 12; - private static final int NUM_THREADS = 4; - - @Test - public void parallelSerializationMatchesSerial() throws Exception { - int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(2)); - - try (CuVSMatrix matrix = new ArrayMatrix(adjacency); - Directory dir = new ByteBuffersDirectory()) { - GPUBuiltHnswGraph serialGraph = newSingleLayerGraph(matrix); - GPUBuiltHnswGraph parallelGraph = newSingleLayerGraph(matrix); - assertSerialAndParallelMatch(serialGraph, parallelGraph, dir); - } - } - - @Test - public void parallelSerializationMatchesSerialAcrossFixedWaveBoundary() throws Exception { - int numNodes = AcceleratedHNSWUtils.SERIALIZATION_WAVE_NODES + 1; - assertTrue(numNodes > AcceleratedHNSWUtils.SERIALIZATION_WAVE_NODES); - - try (Directory dir = new ByteBuffersDirectory()) { - assertSerialAndParallelMatch(new LazyEmptyGraph(numNodes), new LazyEmptyGraph(numNodes), dir); - } - } - - private static void assertSerialAndParallelMatch( - GPUBuiltHnswGraph serialGraph, GPUBuiltHnswGraph parallelGraph, Directory dir) - throws Exception { - int[][] serialOffsets; - try (IndexOutput out = dir.createOutput("serial", IOContext.DEFAULT)) { - serialOffsets = AcceleratedHNSWUtils.writeGraph(serialGraph, out); - } - int[][] parallelOffsets; - try (IndexOutput out = dir.createOutput("parallel", IOContext.DEFAULT)) { - parallelOffsets = AcceleratedHNSWUtils.writeGraph(parallelGraph, out, NUM_THREADS); - } - - assertEquals(serialOffsets.length, parallelOffsets.length); - for (int level = 0; level < serialOffsets.length; level++) { - assertArrayEquals(serialOffsets[level], parallelOffsets[level]); - } - assertArrayEquals(readAllBytes(dir, "serial"), readAllBytes(dir, "parallel")); - } - - private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency) { - return new GPUBuiltHnswGraph( - NUM_NODES, /* dimensions= */ 4, Arrays.asList((int[]) null), List.of(layer0Adjacency)); - } - - private static byte[] readAllBytes(Directory dir, String name) throws Exception { - try (IndexInput in = dir.openInput(name, IOContext.DEFAULT)) { - byte[] bytes = new byte[(int) in.length()]; - in.readBytes(bytes, 0, bytes.length); - return bytes; - } - } - - private static int[][] randomAdjacency(int numNodes, int degree, Random random) { - int[][] adjacency = new int[numNodes][degree]; - for (int[] row : adjacency) { - for (int j = 0; j < degree; j++) { - row[j] = random.nextInt(numNodes); - } - } - return adjacency; - } - - /** Supplies an empty graph lazily so the fixed wave boundary can be tested with little heap. */ - private static final class LazyEmptyGraph extends GPUBuiltHnswGraph { - private final int graphSize; - - LazyEmptyGraph(int graphSize) throws IOException { - super( - 0, - /* dimensions= */ 4, - Arrays.asList((int[]) null), - List.of(new ArrayMatrix(new int[0][])), - 1); - this.graphSize = graphSize; - } - - @Override - public int size() { - return graphSize; - } - - @Override - public int maxConn() { - return 0; - } - - @Override - public NodesIterator getNodesOnLevel(int level) { - return new RangeNodesIterator(level == 0 ? graphSize : 0); - } - - @Override - public NeighborArray getNeighbors(int level, int node) { - return null; - } - } - - private static final class RangeNodesIterator extends GPUBuiltHnswGraph.NodesIterator { - private int current = -1; - - RangeNodesIterator(int size) { - super(size); - } - - @Override - public boolean hasNext() { - return current + 1 < size; - } - - @Override - public int nextInt() { - return ++current; - } - - @Override - public int consume(int[] dest) { - int count = Math.min(dest.length, size - (current + 1)); - for (int i = 0; i < count; i++) { - dest[i] = ++current; - } - return count; - } - } - - private static final class ArrayMatrix implements CuVSMatrix { - private final int[][] rows; - - ArrayMatrix(int[][] rows) { - this.rows = rows; - } - - @Override - public long size() { - return rows.length; - } - - @Override - public long columns() { - return rows.length == 0 ? 0 : rows[0].length; - } - - @Override - public DataType dataType() { - return DataType.INT; - } - - @Override - public RowView getRow(long row) { - int[] values = rows[Math.toIntExact(row)]; - return new RowView() { - @Override - public long size() { - return values.length; - } - - @Override - public int getAsInt(long index) { - return values[Math.toIntExact(index)]; - } - - @Override - public float getAsFloat(long index) { - throw new UnsupportedOperationException(); - } - - @Override - public byte getAsByte(long index) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(int[] target) { - System.arraycopy(values, 0, target, 0, values.length); - } - - @Override - public void toArray(float[] target) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(byte[] target) { - throw new UnsupportedOperationException(); - } - }; - } - - @Override - public void toArray(int[][] target) { - for (int i = 0; i < rows.length; i++) { - System.arraycopy(rows[i], 0, target[i], 0, rows[i].length); - } - } - - @Override - public void toArray(float[][] target) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(byte[][] target) { - throw new UnsupportedOperationException(); - } - - @Override - public void toHost(CuVSHostMatrix target) { - throw new UnsupportedOperationException(); - } - - @Override - public CuVSHostMatrix toHost() { - throw new UnsupportedOperationException(); - } - - @Override - public void toDevice(CuVSDeviceMatrix target, CuVSResources resources) { - throw new UnsupportedOperationException(); - } - - @Override - public CuVSDeviceMatrix toDevice(CuVSResources resources) { - throw new UnsupportedOperationException(); - } - - @Override - public void close() {} - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsPersistedIndex.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsPersistedIndex.java deleted file mode 100644 index 6a5fcb32f1..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsPersistedIndex.java +++ /dev/null @@ -1,126 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; -import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; -import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; - -import java.util.Random; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.codecs.KnnVectorsReader; -import org.apache.lucene.codecs.hnsw.HnswGraphProvider; -import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.CodecReader; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.FloatVectorValues; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.LeafReader; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; -import org.apache.lucene.store.Directory; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; -import org.apache.lucene.tests.util.TestUtil; -import org.apache.lucene.util.hnsw.HnswGraph; -import org.junit.Test; - -@SuppressSysoutChecks(bugUrl = "") -public class TestWriterThreadsPersistedIndex extends LuceneTestCase { - - private static final String VECTOR_FIELD = "vector"; - private static final int VECTOR_COUNT = AcceleratedHNSWUtils.PARALLEL_MIN_NODES + 1; - private static final int DIMENSIONS = 32; - - @Test - public void testParallelGraphRoundTripAboveThreshold() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - AcceleratedHNSWParams params = - new AcceleratedHNSWParams.Builder() - .withWriterThreads(4) - .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) - .withIntermediateGraphDegree(32) - .withGraphDegree(16) - .withHNSWLayer(1) - .build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params); - Random random = new Random(0x2594L); - - try (Directory directory = newDirectory()) { - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(VECTOR_COUNT + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); - try (IndexWriter writer = new IndexWriter(directory, config)) { - for (int id = 0; id < VECTOR_COUNT; id++) { - float[] vector = new float[DIMENSIONS]; - for (int dimension = 0; dimension < DIMENSIONS; dimension++) { - vector[dimension] = random.nextFloat(); - } - Document document = new Document(); - document.add(new StringField("id", Integer.toString(id), Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, vector, EUCLIDEAN)); - writer.addDocument(document); - } - } - - TestUtil.checkIndex(directory); - try (DirectoryReader reader = DirectoryReader.open(directory)) { - assertEquals(1, reader.leaves().size()); - assertEquals(VECTOR_COUNT, reader.numDocs()); - LeafReader leaf = getOnlyLeafReader(reader); - HnswGraph graph = graphOf(leaf); - assertEquals(VECTOR_COUNT, graph.size()); - int arcs = 0; - HnswGraph.NodesIterator nodes = graph.getNodesOnLevel(0); - while (nodes.hasNext()) { - int node = nodes.nextInt(); - graph.seek(0, node); - for (int neighbor = graph.nextNeighbor(); - neighbor != NO_MORE_DOCS; - neighbor = graph.nextNeighbor()) { - assertTrue(neighbor >= 0); - assertTrue(neighbor < VECTOR_COUNT); - arcs++; - } - } - assertTrue("persisted graph contains no arcs", arcs > 0); - - int queryNode = graph.entryNode(); - assertTrue(queryNode >= 0); - assertTrue(queryNode < VECTOR_COUNT); - FloatVectorValues values = leaf.getFloatVectorValues(VECTOR_FIELD); - assertNotNull(values); - float[] query = values.vectorValue(queryNode).clone(); - int queryDoc = values.ordToDoc(queryNode); - String queryId = leaf.storedFields().document(queryDoc).get("id"); - - IndexSearcher searcher = new IndexSearcher(reader); - var hits = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, query, 10), 10); - assertEquals(10, hits.scoreDocs.length); - boolean foundQueryNode = false; - for (var hit : hits.scoreDocs) { - foundQueryNode |= queryId.equals(searcher.storedFields().document(hit.doc).get("id")); - } - assertTrue("the entry-node vector must be returned for its own query", foundQueryNode); - } - } - } - - private static HnswGraph graphOf(LeafReader leaf) throws Exception { - KnnVectorsReader reader = ((CodecReader) leaf).getVectorReader(); - if (reader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { - reader = fieldsReader.getFieldReader(VECTOR_FIELD); - } - return ((HnswGraphProvider) reader).getGraph(VECTOR_FIELD); - } -} From 0f37c156a1a55f13008b52f38bef5b823de41fa5 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Wed, 7 Oct 2026 04:13:43 +0000 Subject: [PATCH 16/21] Make temporary graph copy budget configurable --- ...vidia-cuvs-lucene-acceleratedhnswparams.md | 35 ++++++ fern/pages/user_guide/lucene.md | 4 +- .../cuvs/lucene/AcceleratedHNSWParams.java | 39 ++++++ .../cuvs/lucene/AcceleratedHNSWUtils.java | 59 ++++++++- .../nvidia/cuvs/lucene/GPUBuiltHnswGraph.java | 112 +++++++++++++++--- .../cuvs/lucene/GraphCopyMemoryBudget.java | 109 ++++------------- .../cuvs/lucene/GraphProcessingTrace.java | 45 +++++-- .../Lucene99AcceleratedHNSWVectorsWriter.java | 1 + ...ratedHNSWBinaryQuantizedVectorsWriter.java | 1 + ...ratedHNSWScalarQuantizedVectorsWriter.java | 1 + .../lucene/TestAcceleratedHNSWParams.java | 21 ++++ .../lucene/TestGraphCopyMemoryBudget.java | 106 +++++++---------- .../TestParallelGraphMaterialization.java | 40 ++++++- 13 files changed, 394 insertions(+), 179 deletions(-) diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md index f4ca78f79f..5ac9c2b54b 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md @@ -41,6 +41,18 @@ HNSW graph processing threads parameter _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:164`_ +### getGraphCopyMemoryBudgetBytes + +```java +public long getGraphCopyMemoryBudgetBytes() +``` + +Get the configured temporary host-copy budget for parallel device-graph materialization. + +**Returns** + +graph-copy memory budget in bytes + ### getIntermediateGraphDegree ```java @@ -271,6 +283,29 @@ instance of `Builder` _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:366`_ +### withGraphCopyMemoryBudgetBytes + +```java +public Builder withGraphCopyMemoryBudgetBytes(long graphCopyMemoryBudgetBytes) +``` + +Set the per-operation ceiling for the raw temporary host adjacency copy used by parallel +device-graph materialization. Concurrent copies in the same class loader share reservations; +applications that require one ceiling across codecs should configure the same value for each +codec. This setting does not cap the heap-backed Lucene graph and is not a guarantee of physical +memory availability. A value of `0` preserves the serial fallback without making a temporary +device-to-host graph copy. The default is 42 GiB (`42L << 30` bytes). + +**Parameters** + +| Name | Description | +| --- | --- | +| `graphCopyMemoryBudgetBytes` | graph-copy memory budget in bytes | + +**Returns** + +instance of `Builder` + ### withIntermediateGraphDegree ```java diff --git a/fern/pages/user_guide/lucene.md b/fern/pages/user_guide/lucene.md index 670a67305f..be2bb0eb48 100644 --- a/fern/pages/user_guide/lucene.md +++ b/fern/pages/user_guide/lucene.md @@ -231,7 +231,9 @@ Both `AcceleratedHNSWParams` and `GPUSearchParams` default to a `HEURISTIC` stra - For the accelerated HNSW codecs, set `maxConn` and `beamWidth`, the HNSW parameters you would tune on the CPU. cuVS derives graph degrees and the build algorithm from them. These two values also configure the CPU fallback writer, so one setting covers both paths. - For the GPU search codec, set `buildQuality`. Higher values spend more build time for a higher-quality graph. This codec also passes `graphDegree` into the heuristic, so leave it at its default unless you intend to cap the graph. -`writerThreads` controls native cuVS build concurrency. The accelerated HNSW codecs default to one writer thread, while the GPU search codec defaults to 32. Accelerated HNSW also offers `graphThreads` for the later CPU graph-materialization and serialization stages; it defaults to one and counts the calling thread. Parallel materialization of a device-backed graph uses a temporary host copy only when the JVM reports enough free physical memory for that copy, the Lucene graph, and an adjacency-sized safety allowance. Otherwise, materialization remains serial. +`writerThreads` controls native cuVS build concurrency. The accelerated HNSW codecs default to one writer thread, while the GPU search codec defaults to 32. Accelerated HNSW also offers `graphThreads` for the later CPU graph-materialization and serialization stages; it defaults to one and counts the calling thread. Parallel materialization of a device-backed graph requires a temporary host copy of its INT32 adjacency payload. `graphCopyMemoryBudgetBytes` sets the per-operation ceiling for that copy and defaults to 42 GiB (`42L << 30` bytes). Copies in the same class loader reserve against one shared counter; configure the same ceiling on every codec when the application requires one classloader-wide policy. A copy that would exceed the ceiling falls back to serial device-row reads. + +The graph-copy budget controls only the extra temporary host copy used by the parallel path. It does not cap the heap-backed Lucene graph (which both paths build), measure current physical memory, or guarantee that an admitted native allocation will succeed. Long-lived applications such as Solr should choose a value that leaves headroom for Lucene, the JVM heap, filesystem cache, native libraries, and other cores. Setting the value to `0` disables device-graph copying without disabling indexing. Switching either class to the `CUSTOM` strategy exposes the underlying CAGRA parameters directly, including `graphDegree`, `intermediateGraphDegree`, the graph build algorithm, and its parameters. Use `CUSTOM` only when you have measurements that justify specific values; the defaults derived by cuVS are a better starting point. For background on the parameters themselves, see the [CAGRA indexing guide](/user-guide/api-guides/indexing-guide/cagra) and the [tuning guide](/getting-started/introduction/tuning-indexes). diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index 1f8c13609b..870f967a35 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -53,6 +53,7 @@ public static enum Strategy { public static final int DEFAULT_WRITER_THREADS = 1; public static final int DEFAULT_GRAPH_THREADS = 1; + public static final long DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES = 42L << 30; public static final int DEFAULT_INT_GRAPH_DEGREE = 128; public static final int DEFAULT_GRAPH_DEGREE = 64; public static final int DEFAULT_HNSW_LAYERS = 1; @@ -79,6 +80,7 @@ public static enum Strategy { private final int writerThreads; private final int graphThreads; + private final long graphCopyMemoryBudgetBytes; private final int intermediateGraphDegree; private final int graphdegree; private final int hnswLayers; @@ -99,6 +101,8 @@ public static enum Strategy { * @param writerThreads Number of native cuVS writer threads to use. * @param graphThreads Maximum threads per HNSW graph materialization or serialization operation, * including the calling thread. + * @param graphCopyMemoryBudgetBytes Per-operation ceiling for the raw temporary host adjacency + * copy used by parallel device-graph materialization. * @param intermediateGraphDegree The intermediate graph degree while building the CAGRA index. * @param graphdegree The graph degree to use while building the CAGRA index. * @param hnswLayers The number of HNSW layers to build in the HNSW index. @@ -116,6 +120,7 @@ public static enum Strategy { private AcceleratedHNSWParams( int writerThreads, int graphThreads, + long graphCopyMemoryBudgetBytes, int intermediateGraphDegree, int graphdegree, int hnswLayers, @@ -132,6 +137,7 @@ private AcceleratedHNSWParams( super(); this.writerThreads = writerThreads; this.graphThreads = graphThreads; + this.graphCopyMemoryBudgetBytes = graphCopyMemoryBudgetBytes; this.intermediateGraphDegree = intermediateGraphDegree; this.graphdegree = graphdegree; this.hnswLayers = hnswLayers; @@ -166,6 +172,15 @@ public int getGraphThreads() { return graphThreads; } + /** + * Get the configured temporary host-copy budget for parallel device-graph materialization. + * + * @return graph-copy memory budget in bytes + */ + public long getGraphCopyMemoryBudgetBytes() { + return graphCopyMemoryBudgetBytes; + } + /** * Get the intermediate graph degree * @@ -293,6 +308,8 @@ public String toString() { + writerThreads + ", graphThreads=" + graphThreads + + ", graphCopyMemoryBudgetBytes=" + + graphCopyMemoryBudgetBytes + ", intermediateGraphDegree=" + intermediateGraphDegree + ", graphdegree=" @@ -329,6 +346,7 @@ public static class Builder { private int writerThreads = DEFAULT_WRITER_THREADS; private int graphThreads = DEFAULT_GRAPH_THREADS; + private long graphCopyMemoryBudgetBytes = DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES; private int intermediateGraphDegree = DEFAULT_INT_GRAPH_DEGREE; private int graphdegree = DEFAULT_GRAPH_DEGREE; private int hnswLayers = DEFAULT_HNSW_LAYERS; @@ -369,6 +387,23 @@ public Builder withGraphThreads(int graphThreads) { return this; } + /** + * Set the per-operation ceiling for the raw temporary host adjacency copy used by parallel + * device-graph materialization. Concurrent copies in the same class loader share reservations; + * applications that require one ceiling across codecs should configure the same value for each + * codec. This setting does not cap the heap-backed Lucene graph, which is allocated by both the + * serial and parallel paths, and it is not a guarantee of physical memory availability. A value + * of {@code 0} disables the temporary copy while preserving the serial fallback. Default value + * - {@value DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES} bytes. + * + * @param graphCopyMemoryBudgetBytes graph-copy memory budget in bytes + * @return instance of {@link Builder} + */ + public Builder withGraphCopyMemoryBudgetBytes(long graphCopyMemoryBudgetBytes) { + this.graphCopyMemoryBudgetBytes = graphCopyMemoryBudgetBytes; + return this; + } + /** * Set the intermediate graph degree to use while building CAGRA index * Valid range - Minimum: {@value MIN_INT_GRAPH_DEG}, Maximum: {@value MAX_INT_GRAPH_DEG} @@ -548,6 +583,9 @@ private void validate() throws IllegalArgumentException { "writerThreads", writerThreads, MIN_WRITER_THREADS, MAX_WRITER_THREADS); ParameterValidation.checkRange( "graphThreads", graphThreads, MIN_GRAPH_THREADS, MAX_GRAPH_THREADS); + if (graphCopyMemoryBudgetBytes < 0) { + throw new IllegalArgumentException("graphCopyMemoryBudgetBytes must be non-negative."); + } ParameterValidation.checkRange( "intermediateGraphDegree", intermediateGraphDegree, MIN_INT_GRAPH_DEG, MAX_INT_GRAPH_DEG); ParameterValidation.checkRange("graphdegree", graphdegree, MIN_GRAPH_DEG, MAX_GRAPH_DEG); @@ -591,6 +629,7 @@ public AcceleratedHNSWParams build() { return new AcceleratedHNSWParams( writerThreads, graphThreads, + graphCopyMemoryBudgetBytes, intermediateGraphDegree, graphdegree, hnswLayers, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 5257af0509..3d27e7f717 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -114,6 +114,7 @@ private static GPUBuiltHnswGraph createMultiLayerHnswGraph( params, quantization, graphThreads, + AcceleratedHNSWParams.DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES, GraphProcessingTrace.disabled()); } @@ -128,6 +129,31 @@ private static GPUBuiltHnswGraph createMultiLayerHnswGraph( int graphThreads, GraphProcessingTrace graphProcessingTrace) throws Throwable { + return createMultiLayerHnswGraph( + size, + dimensions, + adjacencyListMatrix, + vectors, + hnswLayers, + params, + quantization, + graphThreads, + AcceleratedHNSWParams.DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES, + graphProcessingTrace); + } + + private static GPUBuiltHnswGraph createMultiLayerHnswGraph( + int size, + int dimensions, + CuVSMatrix adjacencyListMatrix, + List vectors, + int hnswLayers, + CagraIndexParams params, + QuantizationType quantization, + int graphThreads, + long graphCopyMemoryBudgetBytes, + GraphProcessingTrace graphProcessingTrace) + throws Throwable { int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); @@ -216,7 +242,13 @@ private static GPUBuiltHnswGraph createMultiLayerHnswGraph( // The graph eagerly copies all adjacency rows, so generated upper matrices can now close. return new GPUBuiltHnswGraph( - size, dimensions, layerNodes, layerAdjacencies, graphThreads, graphProcessingTrace); + size, + dimensions, + layerNodes, + layerAdjacencies, + graphThreads, + graphCopyMemoryBudgetBytes, + graphProcessingTrace); } catch (Throwable t) { failure = t; throw t; @@ -272,6 +304,7 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( params, quantization, graphThreads, + AcceleratedHNSWParams.DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES, GraphProcessingTrace.disabled()); } @@ -285,6 +318,29 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( int graphThreads, GraphProcessingTrace graphProcessingTrace) throws Throwable { + return createMultiLayerHnswGraph( + dimensions, + adjacencyListMatrix, + vectorDataset, + hnswLayers, + params, + quantization, + graphThreads, + AcceleratedHNSWParams.DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES, + graphProcessingTrace); + } + + static GPUBuiltHnswGraph createMultiLayerHnswGraph( + int dimensions, + CuVSMatrix adjacencyListMatrix, + CuVSMatrix vectorDataset, + int hnswLayers, + CagraIndexParams params, + QuantizationType quantization, + int graphThreads, + long graphCopyMemoryBudgetBytes, + GraphProcessingTrace graphProcessingTrace) + throws Throwable { int size = Math.toIntExact(vectorDataset.size()); // Matrix columns are the stored width: binary vectors are bit-packed, while scalar and float // vectors store one value per dimension. @@ -318,6 +374,7 @@ public int size() { params, quantization, graphThreads, + graphCopyMemoryBudgetBytes, graphProcessingTrace); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index 61c73b10ab..2ac812771d 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -71,7 +71,12 @@ public GPUBuiltHnswGraph( size, dimensions, materialize( - size, layerNodes, layerAdjacencies, graphThreads, GraphProcessingTrace.disabled())); + size, + layerNodes, + layerAdjacencies, + graphThreads, + AcceleratedHNSWParams.DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES, + GraphProcessingTrace.disabled())); } GPUBuiltHnswGraph( @@ -85,7 +90,32 @@ public GPUBuiltHnswGraph( this( size, dimensions, - materialize(size, layerNodes, layerAdjacencies, graphThreads, graphProcessingTrace)); + layerNodes, + layerAdjacencies, + graphThreads, + AcceleratedHNSWParams.DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES, + graphProcessingTrace); + } + + GPUBuiltHnswGraph( + int size, + int dimensions, + List layerNodes, + List layerAdjacencies, + int graphThreads, + long graphCopyMemoryBudgetBytes, + GraphProcessingTrace graphProcessingTrace) + throws IOException { + this( + size, + dimensions, + materialize( + size, + layerNodes, + layerAdjacencies, + graphThreads, + graphCopyMemoryBudgetBytes, + graphProcessingTrace)); } private GPUBuiltHnswGraph(int size, int dimensions, MaterializedGraph graph) { @@ -100,7 +130,13 @@ private GPUBuiltHnswGraph(int size, int dimensions, MaterializedGraph graph) { private static MaterializedGraph materializeSerial( int size, List layerNodes, List layerAdjacencies) { try { - return materialize(size, layerNodes, layerAdjacencies, 1, GraphProcessingTrace.disabled()); + return materialize( + size, + layerNodes, + layerAdjacencies, + 1, + AcceleratedHNSWParams.DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES, + GraphProcessingTrace.disabled()); } catch (IOException impossible) { throw new AssertionError( "serial graph materialization cannot fail with IOException", impossible); @@ -112,19 +148,29 @@ private static MaterializedGraph materialize( List layerNodes, List layerAdjacencies, int graphThreads, + long graphCopyMemoryBudgetBytes, GraphProcessingTrace graphProcessingTrace) throws IOException { List upperLayerNodes = new ArrayList<>(); List upperLayerNeighbors = new ArrayList<>(); NeighborArray[] baseLayerNeighbors = - fillNeighborArray(layerAdjacencies.get(0), size, graphThreads, graphProcessingTrace); + fillNeighborArray( + layerAdjacencies.get(0), + size, + graphThreads, + graphCopyMemoryBudgetBytes, + graphProcessingTrace); for (int level = 1; level < layerAdjacencies.size(); level++) { int[] nodes = layerNodes.get(level); upperLayerNodes.add(nodes); upperLayerNeighbors.add( fillNeighborArray( - layerAdjacencies.get(level), nodes.length, graphThreads, graphProcessingTrace)); + layerAdjacencies.get(level), + nodes.length, + graphThreads, + graphCopyMemoryBudgetBytes, + graphProcessingTrace)); } return new MaterializedGraph( layerAdjacencies.size(), upperLayerNodes, baseLayerNeighbors, upperLayerNeighbors); @@ -140,7 +186,11 @@ private static MaterializedGraph materialize( * @return the NeighborArray */ private static NeighborArray[] fillNeighborArray( - CuVSMatrix adjacency, int size, int graphThreads, GraphProcessingTrace graphProcessingTrace) + CuVSMatrix adjacency, + int size, + int graphThreads, + long graphCopyMemoryBudgetBytes, + GraphProcessingTrace graphProcessingTrace) throws IOException { if (graphThreads <= 1 || size < PARALLEL_MIN_NODES) { NeighborArray[] neighbors = fillNeighborArraySerial(adjacency, size); @@ -160,7 +210,8 @@ private static NeighborArray[] fillNeighborArray( deviceAdjacency, size, graphThreads, - GraphCopyMemoryBudget.system(), + GraphCopyMemoryBudget.shared(), + graphCopyMemoryBudgetBytes, () -> newHostMatrix(deviceAdjacency), graphProcessingTrace); } @@ -189,7 +240,31 @@ static NeighborArray[] materializeDeviceAdjacency( Supplier hostCopyFactory) throws IOException { return materializeDeviceAdjacency( - source, size, graphThreads, memoryBudget, hostCopyFactory, GraphProcessingTrace.disabled()); + source, + size, + graphThreads, + memoryBudget, + AcceleratedHNSWParams.DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES, + hostCopyFactory, + GraphProcessingTrace.disabled()); + } + + static NeighborArray[] materializeDeviceAdjacency( + CuVSDeviceMatrix source, + int size, + int graphThreads, + GraphCopyMemoryBudget memoryBudget, + long graphCopyMemoryBudgetBytes, + Supplier hostCopyFactory) + throws IOException { + return materializeDeviceAdjacency( + source, + size, + graphThreads, + memoryBudget, + graphCopyMemoryBudgetBytes, + hostCopyFactory, + GraphProcessingTrace.disabled()); } private static NeighborArray[] materializeDeviceAdjacency( @@ -197,30 +272,37 @@ private static NeighborArray[] materializeDeviceAdjacency( int size, int graphThreads, GraphCopyMemoryBudget memoryBudget, + long graphCopyMemoryBudgetBytes, Supplier hostCopyFactory, GraphProcessingTrace graphProcessingTrace) throws IOException { Optional reservation = - memoryBudget.tryReserve(source.size(), source.columns()); + memoryBudget.tryReserve(source.size(), source.columns(), graphCopyMemoryBudgetBytes); + long requiredCopyBytes = + GraphCopyMemoryBudget.requiredCopyBytes(source.size(), source.columns()); if (reservation.isEmpty()) { NeighborArray[] neighbors = fillNeighborArraySerial(source, size); - graphProcessingTrace.record( - GraphProcessingTrace.Stage.MATERIALIZATION, + graphProcessingTrace.recordCopyAdmission( GraphProcessingTrace.Mode.SERIAL, GraphProcessingTrace.Reason.MEMORY_ADMISSION_DENIED, graphThreads, - size); + size, + source.columns(), + requiredCopyBytes, + graphCopyMemoryBudgetBytes); return neighbors; } try (GraphCopyMemoryBudget.Reservation ignored = reservation.orElseThrow(); CuVSHostMatrix hostCopy = copyToHost(source, hostCopyFactory)) { NeighborArray[] neighbors = fillNeighborArrayParallel(hostCopy, size, graphThreads); - graphProcessingTrace.record( - GraphProcessingTrace.Stage.MATERIALIZATION, + graphProcessingTrace.recordCopyAdmission( GraphProcessingTrace.Mode.PARALLEL, GraphProcessingTrace.Reason.DEVICE_HOST_COPY, graphThreads, - size); + size, + source.columns(), + requiredCopyBytes, + graphCopyMemoryBudgetBytes); return neighbors; } } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java index 0a03e884d8..67dbbe618f 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java @@ -4,129 +4,66 @@ */ package com.nvidia.cuvs.lucene; -import com.sun.management.OperatingSystemMXBean; -import java.lang.management.ManagementFactory; -import java.util.Objects; import java.util.Optional; -import org.apache.lucene.util.RamUsageEstimator; -import org.apache.lucene.util.hnsw.NeighborArray; /** Coordinates temporary native graph copies across concurrent segment flushes. */ final class GraphCopyMemoryBudget { - private static final long NEIGHBOR_ARRAY_SHALLOW_BYTES = - RamUsageEstimator.shallowSizeOfInstance(NeighborArray.class); + private static final GraphCopyMemoryBudget SHARED = new GraphCopyMemoryBudget(); - private static final GraphCopyMemoryBudget SYSTEM = - new GraphCopyMemoryBudget(GraphCopyMemoryBudget::readSystemMemory); + private long reservedCopyBytes; - private final MemoryProbe memoryProbe; - private long reservedHeadroomBytes; - - GraphCopyMemoryBudget(MemoryProbe memoryProbe) { - this.memoryProbe = Objects.requireNonNull(memoryProbe); - } - - static GraphCopyMemoryBudget system() { - return SYSTEM; + static GraphCopyMemoryBudget shared() { + return SHARED; } /** - * Tries to reserve enough observed free memory for one graph copy and its materialized graph. - * Callers in the same class loader share reservations. This is cooperative admission control, - * not an operating-system memory guarantee. + * Tries to reserve the raw INT32 payload of one temporary device-to-host adjacency copy. + * Reservations are shared by callers in this class loader. Each attempt supplies its own + * ceiling, so applications that require one classloader-wide ceiling must configure the same + * value for every accelerated-HNSW codec in that class loader. */ - synchronized Optional tryReserve(long rows, long columns) { - long requiredHeadroom = requiredHeadroom(rows, columns); - if (requiredHeadroom < 0) { - return Optional.empty(); - } - - MemorySnapshot memory; - try { - memory = memoryProbe.read(); - } catch (RuntimeException unavailable) { + synchronized Optional tryReserve( + long rows, long columns, long configuredBudgetBytes) { + long requiredCopyBytes = requiredCopyBytes(rows, columns); + if (requiredCopyBytes < 0 || configuredBudgetBytes < 0) { return Optional.empty(); } - if (memory == null - || memory.totalBytes() <= 0 - || memory.freeBytes() < 0 - || memory.freeBytes() > memory.totalBytes()) { - return Optional.empty(); - } - - if (reservedHeadroomBytes > memory.freeBytes() - || requiredHeadroom > memory.freeBytes() - reservedHeadroomBytes) { + if (reservedCopyBytes > configuredBudgetBytes + || requiredCopyBytes > configuredBudgetBytes - reservedCopyBytes) { return Optional.empty(); } - reservedHeadroomBytes += requiredHeadroom; - return Optional.of(new Reservation(this, requiredHeadroom)); + reservedCopyBytes += requiredCopyBytes; + return Optional.of(new Reservation(this, requiredCopyBytes)); } - /** - * Estimates peak allocation from the actual matrix shape and current JVM object layout. Besides - * the native host copy and materialized Lucene graph, one adjacency-sized allowance protects - * against allocation races and estimation error while the copy is in flight. - */ - static long requiredHeadroom(long rows, long columns) { + /** Returns the raw INT32 adjacency payload, or {@code -1} for an invalid/overflowing shape. */ + static long requiredCopyBytes(long rows, long columns) { if (rows <= 0 || rows > Integer.MAX_VALUE || columns <= 0 || columns > Integer.MAX_VALUE) { return -1; } try { - long adjacencyBytes = Math.multiplyExact(Math.multiplyExact(rows, columns), Integer.BYTES); - long neighborReferences = arraySize(rows, RamUsageEstimator.NUM_BYTES_OBJECT_REF); - long nodeIds = arraySize(columns, Integer.BYTES); - long scores = arraySize(columns, Float.BYTES); - long bytesPerNode = Math.addExact(NEIGHBOR_ARRAY_SHALLOW_BYTES, nodeIds); - bytesPerNode = Math.addExact(bytesPerNode, scores); - long luceneGraphBytes = - Math.addExact(neighborReferences, Math.multiplyExact(rows, bytesPerNode)); - return Math.addExact(Math.multiplyExact(adjacencyBytes, 2), luceneGraphBytes); + return Math.multiplyExact(Math.multiplyExact(rows, columns), Integer.BYTES); } catch (ArithmeticException overflow) { return -1; } } - private static long arraySize(long length, int bytesPerElement) { - long unaligned = - Math.addExact( - RamUsageEstimator.NUM_BYTES_ARRAY_HEADER, Math.multiplyExact(length, bytesPerElement)); - long alignment = RamUsageEstimator.NUM_BYTES_OBJECT_ALIGNMENT; - long remainder = unaligned % alignment; - return remainder == 0 ? unaligned : Math.addExact(unaligned, alignment - remainder); - } - private synchronized void release(Reservation reservation) { if (reservation.released) { return; } - reservedHeadroomBytes -= reservation.headroomBytes; + reservedCopyBytes -= reservation.copyBytes; reservation.released = true; } - private static MemorySnapshot readSystemMemory() { - java.lang.management.OperatingSystemMXBean platformBean = - ManagementFactory.getOperatingSystemMXBean(); - if (platformBean instanceof OperatingSystemMXBean osBean) { - return new MemorySnapshot(osBean.getTotalMemorySize(), osBean.getFreeMemorySize()); - } - return null; - } - - @FunctionalInterface - interface MemoryProbe { - MemorySnapshot read(); - } - - record MemorySnapshot(long totalBytes, long freeBytes) {} - static final class Reservation implements AutoCloseable { private final GraphCopyMemoryBudget budget; - private final long headroomBytes; + private final long copyBytes; private boolean released; - private Reservation(GraphCopyMemoryBudget budget, long headroomBytes) { + private Reservation(GraphCopyMemoryBudget budget, long copyBytes) { this.budget = budget; - this.headroomBytes = headroomBytes; + this.copyBytes = copyBytes; } @Override diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphProcessingTrace.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphProcessingTrace.java index 54f271fd07..ce33885485 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphProcessingTrace.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphProcessingTrace.java @@ -51,18 +51,43 @@ void record(Stage stage, Mode mode, Reason reason, int requestedThreads, int nod if (infoStream == null || !infoStream.isEnabled(component)) { return; } + infoStream.message(component, message(stage, mode, reason, requestedThreads, nodes)); + } + + void recordCopyAdmission( + Mode mode, + Reason reason, + int requestedThreads, + int nodes, + long columns, + long requiredCopyBytes, + long configuredBudgetBytes) { + if (infoStream == null || !infoStream.isEnabled(component)) { + return; + } infoStream.message( component, - "graph-processing stage=" - + lowerCase(stage) - + " mode=" - + lowerCase(mode) - + " reason=" - + lowerCase(reason) - + " requestedThreads=" - + requestedThreads - + " nodes=" - + nodes); + message(Stage.MATERIALIZATION, mode, reason, requestedThreads, nodes) + + " columns=" + + columns + + " requiredCopyBytes=" + + requiredCopyBytes + + " configuredBudgetBytes=" + + configuredBudgetBytes); + } + + private static String message( + Stage stage, Mode mode, Reason reason, int requestedThreads, int nodes) { + return "graph-processing stage=" + + lowerCase(stage) + + " mode=" + + lowerCase(mode) + + " reason=" + + lowerCase(reason) + + " requestedThreads=" + + requestedThreads + + " nodes=" + + nodes; } private static String lowerCase(Enum value) { diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 8f52eae6e0..4a27d4c473 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -189,6 +189,7 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw params, QuantizationType.NONE, acceleratedHNSWParams.getGraphThreads(), + acceleratedHNSWParams.getGraphCopyMemoryBudgetBytes(), graphProcessingTrace); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 4fa02eb1d5..0ebffd7234 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -187,6 +187,7 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw params, QuantizationType.BINARY, acceleratedHNSWParams.getGraphThreads(), + acceleratedHNSWParams.getGraphCopyMemoryBudgetBytes(), graphProcessingTrace); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 0ad8f24c5e..7a054ed4f9 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -214,6 +214,7 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw params, QuantizationType.SCALAR, acceleratedHNSWParams.getGraphThreads(), + acceleratedHNSWParams.getGraphCopyMemoryBudgetBytes(), graphProcessingTrace); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java index 6794ad1f89..748b13c555 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java @@ -8,6 +8,7 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_BEAM_WIDTH; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_CAGRA_GRAPH_BUILD_ALGO; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_CUVS_DISTANCE_TYPE; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_GRAPH_DEGREE; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_GRAPH_THREADS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_HNSW_LAYERS; @@ -58,6 +59,7 @@ public void testAcceleratedHNSWParamsDefaultValues() { AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); assertEquals(DEFAULT_BEAM_WIDTH, params.getBeamWidth()); assertEquals(DEFAULT_GRAPH_DEGREE, params.getGraphdegree()); + assertEquals(DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES, params.getGraphCopyMemoryBudgetBytes()); assertEquals(DEFAULT_GRAPH_THREADS, params.getGraphThreads()); assertEquals(DEFAULT_HNSW_LAYERS, params.getHnswLayers()); assertEquals(DEFAULT_INT_GRAPH_DEGREE, params.getIntermediateGraphDegree()); @@ -225,6 +227,25 @@ public void testAcceleratedHNSWParamsInvalidGraphThreads() { } } + @Test + public void testAcceleratedHNSWParamsGraphCopyMemoryBudget() { + assertEquals( + 0L, + new AcceleratedHNSWParams.Builder() + .withGraphCopyMemoryBudgetBytes(0) + .build() + .getGraphCopyMemoryBudgetBytes()); + assertEquals( + 80L << 30, + new AcceleratedHNSWParams.Builder() + .withGraphCopyMemoryBudgetBytes(80L << 30) + .build() + .getGraphCopyMemoryBudgetBytes()); + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withGraphCopyMemoryBudgetBytes(-1).build()); + } + @Test public void testWriterAndGraphThreadsAreIndependent() { AcceleratedHNSWParams params = diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java index 6d0d309d3f..c6390935b4 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java @@ -14,46 +14,53 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.util.RamUsageEstimator; -import org.apache.lucene.util.hnsw.NeighborArray; import org.junit.Test; -/** Behavioral specifications for native graph-copy admission control. */ +/** Behavioral specifications for temporary native graph-copy admission control. */ public class TestGraphCopyMemoryBudget extends LuceneTestCase { private static final long TIMEOUT_SECONDS = 10; @Test - public void reservationsFollowEstimatedPeakAcrossSupportedDegrees() { + public void defaultBudgetAdmitsExpectedBenchmarkShapes() { + long defaultBudget = AcceleratedHNSWParams.DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES; + long deep100MRequired = GraphCopyMemoryBudget.requiredCopyBytes(100_000_000L, 32); + long jasper10MRequired = GraphCopyMemoryBudget.requiredCopyBytes(10_000_000L, 32); + + assertEquals(12_800_000_000L, deep100MRequired); + assertEquals(1_280_000_000L, jasper10MRequired); + assertTrue(deep100MRequired < defaultBudget); + assertTrue(jasper10MRequired < defaultBudget); + + GraphCopyMemoryBudget budget = new GraphCopyMemoryBudget(); + try (GraphCopyMemoryBudget.Reservation ignored = + budget.tryReserve(100_000_000L, 32, defaultBudget).orElseThrow()) { + // The default admits one 100M-by-32 temporary copy. + } + } + + @Test + public void reservationsFollowRawInt32PayloadAcrossSupportedDegrees() { int rows = 100; for (int degree : new int[] {1, 32, 512}) { - long required = GraphCopyMemoryBudget.requiredHeadroom(rows, degree); - long adjacencyBytes = (long) rows * degree * Integer.BYTES; - long expected = - 2 * adjacencyBytes - + RamUsageEstimator.shallowSizeOf(new NeighborArray[rows]) - + rows - * (RamUsageEstimator.shallowSizeOfInstance(NeighborArray.class) - + RamUsageEstimator.sizeOf(new int[degree]) - + RamUsageEstimator.sizeOf(new float[degree])); - assertEquals(expected, required); - assertTrue("object layout must be included", required > 4 * adjacencyBytes); + long required = (long) rows * degree * Integer.BYTES; + assertEquals(required, GraphCopyMemoryBudget.requiredCopyBytes(rows, degree)); - GraphCopyMemoryBudget exactBudget = budgetWith(required, required); - try (GraphCopyMemoryBudget.Reservation ignored = reserve(exactBudget, rows, degree)) { - assertTrue(exactBudget.tryReserve(1, 1).isEmpty()); + GraphCopyMemoryBudget exactBudget = new GraphCopyMemoryBudget(); + try (GraphCopyMemoryBudget.Reservation ignored = + exactBudget.tryReserve(rows, degree, required).orElseThrow()) { + assertTrue(exactBudget.tryReserve(1, 1, required).isEmpty()); } - GraphCopyMemoryBudget insufficientBudget = budgetWith(required, required - 1); - assertTrue(insufficientBudget.tryReserve(rows, degree).isEmpty()); + assertTrue(new GraphCopyMemoryBudget().tryReserve(rows, degree, required - 1).isEmpty()); } } @Test - public void concurrentReservationsCannotExceedSharedHeadroom() throws Exception { + public void concurrentReservationsCannotExceedRequestCeiling() throws Exception { long rows = 100; long degree = 16; - long reservationBytes = GraphCopyMemoryBudget.requiredHeadroom(rows, degree); - GraphCopyMemoryBudget budget = budgetWith(2 * reservationBytes, 2 * reservationBytes); + long reservationBytes = GraphCopyMemoryBudget.requiredCopyBytes(rows, degree); + GraphCopyMemoryBudget budget = new GraphCopyMemoryBudget(); int callers = 8; ExecutorService executor = Executors.newFixedThreadPool(callers); CountDownLatch start = new CountDownLatch(1); @@ -68,7 +75,7 @@ public void concurrentReservationsCannotExceedSharedHeadroom() throws Exception () -> { assertTrue(start.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); Optional reservation = - budget.tryReserve(rows, degree); + budget.tryReserve(rows, degree, 2 * reservationBytes); reservation.ifPresent(ignored -> granted.incrementAndGet()); attempted.countDown(); if (reservation.isPresent()) { @@ -95,9 +102,10 @@ public void concurrentReservationsCannotExceedSharedHeadroom() throws Exception @Test public void reservationIsReleasedOnFailureAndCloseIsIdempotent() { - long required = GraphCopyMemoryBudget.requiredHeadroom(100, 16); - GraphCopyMemoryBudget budget = budgetWith(required, required); - GraphCopyMemoryBudget.Reservation failedOperation = reserve(budget, 100, 16); + long required = GraphCopyMemoryBudget.requiredCopyBytes(100, 16); + GraphCopyMemoryBudget budget = new GraphCopyMemoryBudget(); + GraphCopyMemoryBudget.Reservation failedOperation = + budget.tryReserve(100, 16, required).orElseThrow(); RuntimeException expected = new RuntimeException("expected"); RuntimeException actual = @@ -111,42 +119,20 @@ public void reservationIsReleasedOnFailureAndCloseIsIdempotent() { assertSame(expected, actual); failedOperation.close(); - try (GraphCopyMemoryBudget.Reservation replacement = reserve(budget, 100, 16)) { - assertTrue(budget.tryReserve(1, 1).isEmpty()); + try (GraphCopyMemoryBudget.Reservation ignored = + budget.tryReserve(100, 16, required).orElseThrow()) { + assertTrue(budget.tryReserve(1, 1, required).isEmpty()); } } @Test - public void invalidOrUnavailableMemoryInformationFailsClosed() { - assertRejected(() -> null); - assertRejected(() -> new GraphCopyMemoryBudget.MemorySnapshot(0, 0)); - assertRejected(() -> new GraphCopyMemoryBudget.MemorySnapshot(1_000, -1)); - assertRejected(() -> new GraphCopyMemoryBudget.MemorySnapshot(1_000, 1_001)); - assertRejected( - () -> { - throw new UnsupportedOperationException("unavailable"); - }); - - GraphCopyMemoryBudget budget = - budgetWith(/* totalBytes= */ Long.MAX_VALUE, /* freeBytes= */ Long.MAX_VALUE); - assertTrue(budget.tryReserve(Integer.MAX_VALUE, Integer.MAX_VALUE).isEmpty()); - assertTrue(budget.tryReserve(0, 1).isEmpty()); - assertTrue(budget.tryReserve(1, 0).isEmpty()); - assertTrue(budget.tryReserve(-1, 1).isEmpty()); - assertTrue(budget.tryReserve(1, -1).isEmpty()); - } - - private static GraphCopyMemoryBudget budgetWith(long totalBytes, long freeBytes) { - return new GraphCopyMemoryBudget( - () -> new GraphCopyMemoryBudget.MemorySnapshot(totalBytes, freeBytes)); - } - - private static GraphCopyMemoryBudget.Reservation reserve( - GraphCopyMemoryBudget budget, long rows, long columns) { - return budget.tryReserve(rows, columns).orElseThrow(); - } - - private static void assertRejected(GraphCopyMemoryBudget.MemoryProbe probe) { - assertTrue(new GraphCopyMemoryBudget(probe).tryReserve(1, 1).isEmpty()); + public void invalidShapesAndBudgetsFailClosed() { + GraphCopyMemoryBudget budget = new GraphCopyMemoryBudget(); + assertTrue(budget.tryReserve(Integer.MAX_VALUE, Integer.MAX_VALUE, Long.MAX_VALUE).isEmpty()); + assertTrue(budget.tryReserve(0, 1, Long.MAX_VALUE).isEmpty()); + assertTrue(budget.tryReserve(1, 0, Long.MAX_VALUE).isEmpty()); + assertTrue(budget.tryReserve(-1, 1, Long.MAX_VALUE).isEmpty()); + assertTrue(budget.tryReserve(1, -1, Long.MAX_VALUE).isEmpty()); + assertTrue(budget.tryReserve(1, 1, -1).isEmpty()); } } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java index 4a6034b210..06153c1793 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java @@ -97,14 +97,12 @@ public void toHost(CuVSHostMatrix target) { copyCount.incrementAndGet(); } }; - long requiredHeadroom = GraphCopyMemoryBudget.requiredHeadroom(NUM_NODES, DEGREE); - GraphCopyMemoryBudget budget = - new GraphCopyMemoryBudget( - () -> new GraphCopyMemoryBudget.MemorySnapshot(requiredHeadroom, requiredHeadroom)); + long requiredCopyBytes = GraphCopyMemoryBudget.requiredCopyBytes(NUM_NODES, DEGREE); + GraphCopyMemoryBudget budget = new GraphCopyMemoryBudget(); NeighborArray[] neighbors = GPUBuiltHnswGraph.materializeDeviceAdjacency( - source, NUM_NODES, GRAPH_THREADS, budget, () -> hostCopy); + source, NUM_NODES, GRAPH_THREADS, budget, requiredCopyBytes, () -> hostCopy); for (int node = 0; node < NUM_NODES; node++) { assertArrayEquals( @@ -114,11 +112,41 @@ public void toHost(CuVSHostMatrix target) { assertEquals(1, hostCloseCount.get()); assertTrue(executionProbe.threadCount() > 1); try (GraphCopyMemoryBudget.Reservation ignored = - budget.tryReserve(NUM_NODES, DEGREE).orElseThrow()) { + budget.tryReserve(NUM_NODES, DEGREE, requiredCopyBytes).orElseThrow()) { // The first reservation was released after materialization. } } + @Test + public void configuredBudgetControlsWhetherDeviceCopyRuns() throws Exception { + int[][] sourceRows = IntGraphTestMatrix.randomRows(NUM_NODES, DEGREE, 5); + AtomicInteger copyCount = new AtomicInteger(); + CuVSDeviceMatrix source = + new IntGraphTestMatrix.DeviceMatrix(sourceRows, DEGREE) { + @Override + public void toHost(CuVSHostMatrix target) { + copyCount.incrementAndGet(); + } + }; + long requiredCopyBytes = GraphCopyMemoryBudget.requiredCopyBytes(NUM_NODES, DEGREE); + GraphCopyMemoryBudget budget = new GraphCopyMemoryBudget(); + + GPUBuiltHnswGraph.materializeDeviceAdjacency( + source, + NUM_NODES, + GRAPH_THREADS, + budget, + requiredCopyBytes - 1, + () -> new IntGraphTestMatrix.TrackingHostMatrix(new AtomicInteger(), null)); + assertEquals(0, copyCount.get()); + + CuVSHostMatrix hostCopy = + new IntGraphTestMatrix.TrackingHostMatrix(sourceRows, new AtomicInteger(), null, null); + GPUBuiltHnswGraph.materializeDeviceAdjacency( + source, NUM_NODES, GRAPH_THREADS, budget, requiredCopyBytes, () -> hostCopy); + assertEquals(1, copyCount.get()); + } + private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, int numThreads) throws IOException { return new GPUBuiltHnswGraph( From a5ed4e27985c3962063f33e4922516108e652a54 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Wed, 7 Oct 2026 04:30:57 +0000 Subject: [PATCH 17/21] Stabilize overlapping graph copy budgets --- ...vidia-cuvs-lucene-acceleratedhnswparams.md | 79 ++++++++------- fern/pages/user_guide/lucene.md | 4 +- .../cuvs/lucene/AcceleratedHNSWParams.java | 15 +-- .../cuvs/lucene/GraphCopyMemoryBudget.java | 20 +++- .../cuvs/lucene/IntGraphTestMatrix.java | 2 +- .../lucene/TestGraphCopyMemoryBudget.java | 24 +++++ .../TestParallelGraphMaterialization.java | 97 +++++++++++++++++++ 7 files changed, 191 insertions(+), 50 deletions(-) diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md index 5ac9c2b54b..b72df83ebe 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md @@ -24,7 +24,7 @@ Get the native cuVS writer threads parameter. cuVS writer threads parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:154`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:161`_ ### getGraphThreads @@ -39,7 +39,7 @@ includes the calling thread; shared helper capacity may reduce actual concurrenc HNSW graph processing threads parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:164`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:171`_ ### getGraphCopyMemoryBudgetBytes @@ -53,6 +53,8 @@ Get the configured temporary host-copy budget for parallel device-graph material graph-copy memory budget in bytes +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:180`_ + ### getIntermediateGraphDegree ```java @@ -65,7 +67,7 @@ Get the intermediate graph degree the graph degree parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:173`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:189`_ ### getGraphdegree @@ -79,7 +81,7 @@ Get the graph degree the graph degree parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:182`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:198`_ ### getHnswLayers @@ -93,7 +95,7 @@ Get the number of HNSW layers the number of HNSW layers -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:191`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:207`_ ### getMaxConn @@ -107,7 +109,7 @@ Get the max connection parameter the max connection parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:200`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:216`_ ### getBeamWidth @@ -121,7 +123,7 @@ Get the beam width parameter the beam width parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:209`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:225`_ ### getCagraGraphBuildAlgo @@ -135,7 +137,7 @@ Get the CAGRA graph build algorithm the CAGRA graph build algorithm -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:218`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:234`_ ### getCuVSIvfPqParams @@ -149,7 +151,7 @@ Get the instance of `CuVSIvfPqParams` the instance of `CuVSIvfPqParams` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:227`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:243`_ ### getNumMergeWorkers @@ -163,7 +165,7 @@ Get the number of merge workers set to be used in the fallback mechanism the number of merge workers -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:236`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:252`_ ### getMergeExec @@ -177,7 +179,7 @@ Get the instance of the `ExecutorService` to be used in the fallback mechanism the instance of the `ExecutorService` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:245`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:261`_ ### getStrategy @@ -194,7 +196,7 @@ When CUSTOM is chosen, the build algorithm and its parameters (either defaults o get the chosen `Strategy` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:257`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:273`_ ### getCuvsDistanceType @@ -208,7 +210,7 @@ Get the cuvs distance type the distance type -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:266`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:282`_ ### getNNDescentNumIterations @@ -222,7 +224,7 @@ get the number of Iterations to run if building with NN_DESCENT the number of iterations for NN_DESCENT -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:275`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:291`_ ### getHnswHeuristicType @@ -237,7 +239,7 @@ beamWidth. Only consulted under the `Strategy#HEURISTIC` strategy. the `HnswHeuristicType` to hand to cuVS -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:285`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:301`_ ### withWriterThreads @@ -259,7 +261,7 @@ Default value - \{@value DEFAULT_WRITER_THREADS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:353`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:372`_ ### withGraphThreads @@ -281,7 +283,7 @@ includes the calling thread. Valid range - Minimum: \{@value MIN_GRAPH_THREADS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:366`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:385`_ ### withGraphCopyMemoryBudgetBytes @@ -289,12 +291,13 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWP public Builder withGraphCopyMemoryBudgetBytes(long graphCopyMemoryBudgetBytes) ``` -Set the per-operation ceiling for the raw temporary host adjacency copy used by parallel -device-graph materialization. Concurrent copies in the same class loader share reservations; -applications that require one ceiling across codecs should configure the same value for each -codec. This setting does not cap the heap-backed Lucene graph and is not a guarantee of physical -memory availability. A value of `0` preserves the serial fallback without making a temporary -device-to-host graph copy. The default is 42 GiB (`42L << 30` bytes). +Set the ceiling for the raw temporary host adjacency copy used by parallel device-graph +materialization. Copies with the same ceiling share aggregate reservations in this class loader. +A differently configured copy cannot overlap active reservations and uses the serial fallback, so +one codec cannot silently raise another codec's active aggregate ceiling. This setting does not +cap the heap-backed Lucene graph, guarantee physical memory availability, or impose a process-wide +limit. A value of `0` preserves the serial fallback without making a temporary device-to-host graph +copy. The default is 42 GiB (`42L << 30` bytes). **Parameters** @@ -306,6 +309,8 @@ device-to-host graph copy. The default is 42 GiB (`42L << 30` bytes). instance of `Builder` +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:403`_ + ### withIntermediateGraphDegree ```java @@ -326,7 +331,7 @@ Default value - \{@value DEFAULT_INT_GRAPH_DEGREE\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:379`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:416`_ ### withGraphDegree @@ -348,7 +353,7 @@ Default value - \{@value DEFAULT_GRAPH_DEGREE\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:392`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:429`_ ### withHNSWLayer @@ -370,7 +375,7 @@ Default value - \{@value DEFAULT_HNSW_LAYERS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:405`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:442`_ ### withMaxConn @@ -392,7 +397,7 @@ Default value - \{@value DEFAULT_MAX_CONN\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:418`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:455`_ ### withBeamWidth @@ -414,7 +419,7 @@ Default value - \{@value DEFAULT_BEAM_WIDTH\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:431`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:468`_ ### withCagraGraphBuildAlgo @@ -435,7 +440,7 @@ Default value - NN_DESCENT instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:443`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:480`_ ### withCuVSIvfPqParams @@ -455,7 +460,7 @@ Set the instance of `CuVSIvfPqParams` instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:454`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:491`_ ### withNumMergeWorkers @@ -476,7 +481,7 @@ Default value - \{@value DEFAULT_NUM_MERGE_WORKERS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:466`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:503`_ ### withMergeExecutorService @@ -497,7 +502,7 @@ Default value an instance with one thread instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:478`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:515`_ ### withStrategy @@ -523,7 +528,7 @@ Default value - HEURISTIC instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:495`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:532`_ ### withCuvsDistanceType @@ -543,7 +548,7 @@ Set the CuvsDistanceType instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:506`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:543`_ ### withNNDescentNumIterations @@ -566,7 +571,7 @@ Default value - \{@value DEFAULT_NN_DESCENT_NUM_ITERATIONS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:520`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:557`_ ### withHnswHeuristicType @@ -590,7 +595,7 @@ the equivalent HNSW graph (graph degree = 2 * maxConn). instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:535`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:572`_ ### build @@ -604,6 +609,6 @@ Create an instance of `AcceleratedHNSWParams` instance of `AcceleratedHNSWParams` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:582`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:622`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:17`_ diff --git a/fern/pages/user_guide/lucene.md b/fern/pages/user_guide/lucene.md index be2bb0eb48..51b3143ebf 100644 --- a/fern/pages/user_guide/lucene.md +++ b/fern/pages/user_guide/lucene.md @@ -231,9 +231,9 @@ Both `AcceleratedHNSWParams` and `GPUSearchParams` default to a `HEURISTIC` stra - For the accelerated HNSW codecs, set `maxConn` and `beamWidth`, the HNSW parameters you would tune on the CPU. cuVS derives graph degrees and the build algorithm from them. These two values also configure the CPU fallback writer, so one setting covers both paths. - For the GPU search codec, set `buildQuality`. Higher values spend more build time for a higher-quality graph. This codec also passes `graphDegree` into the heuristic, so leave it at its default unless you intend to cap the graph. -`writerThreads` controls native cuVS build concurrency. The accelerated HNSW codecs default to one writer thread, while the GPU search codec defaults to 32. Accelerated HNSW also offers `graphThreads` for the later CPU graph-materialization and serialization stages; it defaults to one and counts the calling thread. Parallel materialization of a device-backed graph requires a temporary host copy of its INT32 adjacency payload. `graphCopyMemoryBudgetBytes` sets the per-operation ceiling for that copy and defaults to 42 GiB (`42L << 30` bytes). Copies in the same class loader reserve against one shared counter; configure the same ceiling on every codec when the application requires one classloader-wide policy. A copy that would exceed the ceiling falls back to serial device-row reads. +`writerThreads` controls native cuVS build concurrency. The accelerated HNSW codecs default to one writer thread, while the GPU search codec defaults to 32. Accelerated HNSW also offers `graphThreads` for the later CPU graph-materialization and serialization stages; it defaults to one and counts the calling thread. Parallel materialization of a device-backed graph requires a temporary host copy of its INT32 adjacency payload. `graphCopyMemoryBudgetBytes` sets the ceiling for that copy and defaults to 42 GiB (`42L << 30` bytes). Copies with the same ceiling reserve against one shared counter in the codec's class loader. A differently configured copy cannot overlap active reservations and falls back to serial device-row reads, so one codec cannot silently raise another codec's active aggregate ceiling. Configure the same value on every accelerated-HNSW codec when the application requires one consistent classloader-wide policy. -The graph-copy budget controls only the extra temporary host copy used by the parallel path. It does not cap the heap-backed Lucene graph (which both paths build), measure current physical memory, or guarantee that an admitted native allocation will succeed. Long-lived applications such as Solr should choose a value that leaves headroom for Lucene, the JVM heap, filesystem cache, native libraries, and other cores. Setting the value to `0` disables device-graph copying without disabling indexing. +The graph-copy budget controls only the extra temporary host copy used by the parallel path. It does not cap the heap-backed Lucene graph (which both paths build), measure current physical memory, guarantee that an admitted native allocation will succeed, or coordinate copies made by separately loaded copies of this library. Long-lived applications such as Solr, Elasticsearch, and OpenSearch should configure every accelerated-HNSW codec consistently and choose a value that leaves headroom for Lucene, the JVM heap, filesystem cache, native libraries, and other components. Setting the value to `0` disables device-graph copying without disabling indexing. Switching either class to the `CUSTOM` strategy exposes the underlying CAGRA parameters directly, including `graphDegree`, `intermediateGraphDegree`, the graph build algorithm, and its parameters. Use `CUSTOM` only when you have measurements that justify specific values; the defaults derived by cuVS are a better starting point. For background on the parameters themselves, see the [CAGRA indexing guide](/user-guide/api-guides/indexing-guide/cagra) and the [tuning guide](/getting-started/introduction/tuning-indexes). diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index 870f967a35..2b31a361a1 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -388,13 +388,14 @@ public Builder withGraphThreads(int graphThreads) { } /** - * Set the per-operation ceiling for the raw temporary host adjacency copy used by parallel - * device-graph materialization. Concurrent copies in the same class loader share reservations; - * applications that require one ceiling across codecs should configure the same value for each - * codec. This setting does not cap the heap-backed Lucene graph, which is allocated by both the - * serial and parallel paths, and it is not a guarantee of physical memory availability. A value - * of {@code 0} disables the temporary copy while preserving the serial fallback. Default value - * - {@value DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES} bytes. + * Set the ceiling for the raw temporary host adjacency copy used by parallel device-graph + * materialization. Copies with the same ceiling share aggregate reservations in this class + * loader. A copy with a different ceiling cannot overlap those reservations and uses the serial + * fallback instead. This setting does not cap the heap-backed Lucene graph, which is allocated + * by both the serial and parallel paths, and it is not a guarantee of physical memory + * availability or a process-wide limit. A value of {@code 0} disables the temporary copy while + * preserving the serial fallback. Default value - {@value + * DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES} bytes. * * @param graphCopyMemoryBudgetBytes graph-copy memory budget in bytes * @return instance of {@link Builder} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java index 67dbbe618f..9959fd3781 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java @@ -11,6 +11,7 @@ final class GraphCopyMemoryBudget { private static final GraphCopyMemoryBudget SHARED = new GraphCopyMemoryBudget(); private long reservedCopyBytes; + private long activeBudgetBytes = -1; static GraphCopyMemoryBudget shared() { return SHARED; @@ -18,9 +19,10 @@ static GraphCopyMemoryBudget shared() { /** * Tries to reserve the raw INT32 payload of one temporary device-to-host adjacency copy. - * Reservations are shared by callers in this class loader. Each attempt supplies its own - * ceiling, so applications that require one classloader-wide ceiling must configure the same - * value for every accelerated-HNSW codec in that class loader. + * Reservations are shared by callers in this class loader. Overlapping reservations must use + * the same configured ceiling; a caller with a different ceiling is denied until the active + * reservations are released. This keeps one caller from silently raising another caller's + * active aggregate ceiling. */ synchronized Optional tryReserve( long rows, long columns, long configuredBudgetBytes) { @@ -28,10 +30,19 @@ synchronized Optional tryReserve( if (requiredCopyBytes < 0 || configuredBudgetBytes < 0) { return Optional.empty(); } + if (requiredCopyBytes > configuredBudgetBytes) { + return Optional.empty(); + } + if (reservedCopyBytes != 0 && activeBudgetBytes != configuredBudgetBytes) { + return Optional.empty(); + } if (reservedCopyBytes > configuredBudgetBytes || requiredCopyBytes > configuredBudgetBytes - reservedCopyBytes) { return Optional.empty(); } + if (reservedCopyBytes == 0) { + activeBudgetBytes = configuredBudgetBytes; + } reservedCopyBytes += requiredCopyBytes; return Optional.of(new Reservation(this, requiredCopyBytes)); } @@ -53,6 +64,9 @@ private synchronized void release(Reservation reservation) { return; } reservedCopyBytes -= reservation.copyBytes; + if (reservedCopyBytes == 0) { + activeBudgetBytes = -1; + } reservation.released = true; } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/IntGraphTestMatrix.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/IntGraphTestMatrix.java index 73ef7c5a78..47a43210d9 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/IntGraphTestMatrix.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/IntGraphTestMatrix.java @@ -124,7 +124,7 @@ public CuVSHostMatrix toHost() { } } - static final class TrackingHostMatrix extends IntGraphTestMatrix implements CuVSHostMatrix { + static class TrackingHostMatrix extends IntGraphTestMatrix implements CuVSHostMatrix { private final AtomicInteger closeCount; private final RuntimeException closeFailure; private final ParallelExecutionProbe executionProbe; diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java index c6390935b4..171673d5c0 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java @@ -26,6 +26,7 @@ public void defaultBudgetAdmitsExpectedBenchmarkShapes() { long deep100MRequired = GraphCopyMemoryBudget.requiredCopyBytes(100_000_000L, 32); long jasper10MRequired = GraphCopyMemoryBudget.requiredCopyBytes(10_000_000L, 32); + assertEquals(42L << 30, defaultBudget); assertEquals(12_800_000_000L, deep100MRequired); assertEquals(1_280_000_000L, jasper10MRequired); assertTrue(deep100MRequired < defaultBudget); @@ -100,6 +101,29 @@ public void concurrentReservationsCannotExceedRequestCeiling() throws Exception } } + @Test + public void overlappingReservationsCannotMixConfiguredCeilings() { + long smallCopyBytes = GraphCopyMemoryBudget.requiredCopyBytes(100, 1); + long largeCopyBytes = GraphCopyMemoryBudget.requiredCopyBytes(300, 1); + long smallCeiling = 2 * smallCopyBytes; + long largeCeiling = 2 * largeCopyBytes; + GraphCopyMemoryBudget budget = new GraphCopyMemoryBudget(); + + try (GraphCopyMemoryBudget.Reservation ignored = + budget.tryReserve(100, 1, smallCeiling).orElseThrow()) { + assertTrue(budget.tryReserve(300, 1, largeCeiling).isEmpty()); + } + try (GraphCopyMemoryBudget.Reservation ignored = + budget.tryReserve(300, 1, largeCeiling).orElseThrow()) { + assertTrue(budget.tryReserve(100, 1, smallCeiling).isEmpty()); + } + + try (GraphCopyMemoryBudget.Reservation ignored = + budget.tryReserve(100, 1, smallCeiling).orElseThrow()) { + // Releasing all reservations resets the active ceiling. + } + } + @Test public void reservationIsReleasedOnFailureAndCloseIsIdempotent() { long required = GraphCopyMemoryBudget.requiredCopyBytes(100, 16); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java index 06153c1793..d263c9896a 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java @@ -9,6 +9,7 @@ import com.nvidia.cuvs.CuVSDeviceMatrix; import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.RowView; import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; @@ -78,6 +79,95 @@ public void toHost(CuVSHostMatrix target) { assertArrayEquals(new Throwable[] {closeFailure}, thrown.getSuppressed()); } + @Test + public void hostAllocationFailureReleasesCopyReservation() { + int[][] rows = new int[][] {{0}}; + CuVSDeviceMatrix source = new IntGraphTestMatrix.DeviceMatrix(rows, 1); + long requiredCopyBytes = GraphCopyMemoryBudget.requiredCopyBytes(1, 1); + GraphCopyMemoryBudget budget = new GraphCopyMemoryBudget(); + RuntimeException allocationFailure = new RuntimeException("allocation failed"); + + RuntimeException thrown = + assertThrows( + RuntimeException.class, + () -> + GPUBuiltHnswGraph.materializeDeviceAdjacency( + source, + 1, + GRAPH_THREADS, + budget, + requiredCopyBytes, + () -> { + throw allocationFailure; + })); + + assertSame(allocationFailure, thrown); + assertBudgetIsReusable(budget, requiredCopyBytes); + } + + @Test + public void copyFailureClosesHostAndReleasesCopyReservation() { + int[][] rows = new int[][] {{0}}; + RuntimeException copyFailure = new RuntimeException("copy failed"); + AtomicInteger hostCloseCount = new AtomicInteger(); + CuVSDeviceMatrix source = + new IntGraphTestMatrix.DeviceMatrix(rows, 1) { + @Override + public void toHost(CuVSHostMatrix target) { + throw copyFailure; + } + }; + CuVSHostMatrix hostCopy = + new IntGraphTestMatrix.TrackingHostMatrix(rows, hostCloseCount, null, null); + long requiredCopyBytes = GraphCopyMemoryBudget.requiredCopyBytes(1, 1); + GraphCopyMemoryBudget budget = new GraphCopyMemoryBudget(); + + RuntimeException thrown = + assertThrows( + RuntimeException.class, + () -> + GPUBuiltHnswGraph.materializeDeviceAdjacency( + source, 1, GRAPH_THREADS, budget, requiredCopyBytes, () -> hostCopy)); + + assertSame(copyFailure, thrown); + assertEquals(1, hostCloseCount.get()); + assertBudgetIsReusable(budget, requiredCopyBytes); + } + + @Test + public void parallelFillFailureClosesHostAndReleasesCopyReservation() { + int[][] rows = new int[][] {{0}}; + RuntimeException fillFailure = new RuntimeException("fill failed"); + AtomicInteger hostCloseCount = new AtomicInteger(); + CuVSDeviceMatrix source = + new IntGraphTestMatrix.DeviceMatrix(rows, 1) { + @Override + public void toHost(CuVSHostMatrix target) { + // The fake host matrix already contains the copied row. + } + }; + CuVSHostMatrix hostCopy = + new IntGraphTestMatrix.TrackingHostMatrix(rows, hostCloseCount, null, null) { + @Override + public RowView getRow(long row) { + throw fillFailure; + } + }; + long requiredCopyBytes = GraphCopyMemoryBudget.requiredCopyBytes(1, 1); + GraphCopyMemoryBudget budget = new GraphCopyMemoryBudget(); + + RuntimeException thrown = + assertThrows( + RuntimeException.class, + () -> + GPUBuiltHnswGraph.materializeDeviceAdjacency( + source, 1, GRAPH_THREADS, budget, requiredCopyBytes, () -> hostCopy)); + + assertSame(fillFailure, thrown); + assertEquals(1, hostCloseCount.get()); + assertBudgetIsReusable(budget, requiredCopyBytes); + } + @Test public void admittedDeviceCopyIsMaterializedInParallelAndReleased() throws Exception { int[][] expectedRows = IntGraphTestMatrix.randomRows(NUM_NODES, DEGREE, 3); @@ -157,6 +247,13 @@ private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, numThreads); } + private static void assertBudgetIsReusable(GraphCopyMemoryBudget budget, long requiredCopyBytes) { + try (GraphCopyMemoryBudget.Reservation ignored = + budget.tryReserve(1, 1, requiredCopyBytes).orElseThrow()) { + // The failed materialization released its reservation. + } + } + private static void assertGraphsEqual(HnswGraph a, HnswGraph b) throws Exception { assertEquals(a.numLevels(), b.numLevels()); for (int level = 0; level < a.numLevels(); level++) { From 84eff73d8c95a4d971e72d947e297d949cb72a8b Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Wed, 7 Oct 2026 05:02:34 +0000 Subject: [PATCH 18/21] Clarify graph copy budget integration scope --- fern/pages/user_guide/lucene.md | 2 +- .../java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/fern/pages/user_guide/lucene.md b/fern/pages/user_guide/lucene.md index 51b3143ebf..969ad17ec6 100644 --- a/fern/pages/user_guide/lucene.md +++ b/fern/pages/user_guide/lucene.md @@ -233,7 +233,7 @@ Both `AcceleratedHNSWParams` and `GPUSearchParams` default to a `HEURISTIC` stra `writerThreads` controls native cuVS build concurrency. The accelerated HNSW codecs default to one writer thread, while the GPU search codec defaults to 32. Accelerated HNSW also offers `graphThreads` for the later CPU graph-materialization and serialization stages; it defaults to one and counts the calling thread. Parallel materialization of a device-backed graph requires a temporary host copy of its INT32 adjacency payload. `graphCopyMemoryBudgetBytes` sets the ceiling for that copy and defaults to 42 GiB (`42L << 30` bytes). Copies with the same ceiling reserve against one shared counter in the codec's class loader. A differently configured copy cannot overlap active reservations and falls back to serial device-row reads, so one codec cannot silently raise another codec's active aggregate ceiling. Configure the same value on every accelerated-HNSW codec when the application requires one consistent classloader-wide policy. -The graph-copy budget controls only the extra temporary host copy used by the parallel path. It does not cap the heap-backed Lucene graph (which both paths build), measure current physical memory, guarantee that an admitted native allocation will succeed, or coordinate copies made by separately loaded copies of this library. Long-lived applications such as Solr, Elasticsearch, and OpenSearch should configure every accelerated-HNSW codec consistently and choose a value that leaves headroom for Lucene, the JVM heap, filesystem cache, native libraries, and other components. Setting the value to `0` disables device-graph copying without disabling indexing. +The graph-copy budget controls only the extra temporary host copy used by the parallel path. It does not cap the heap-backed Lucene graph (which both paths build), measure current physical memory, guarantee that an admitted native allocation will succeed, or coordinate copies made by separately loaded copies of this library. Long-lived applications such as Solr, Elasticsearch, and OpenSearch should configure every accelerated-HNSW codec consistently and choose a value that leaves headroom for Lucene, the JVM heap, filesystem cache, native libraries, and other components. The core API exposes the setting through `AcceleratedHNSWParams.Builder`; a server integration must map its operator-facing configuration into that builder. Name-only codec SPI construction uses the defaults, including `graphThreads = 1`, for which no temporary graph copy is made. Setting the budget to `0` disables device-graph copying without disabling indexing. Switching either class to the `CUSTOM` strategy exposes the underlying CAGRA parameters directly, including `graphDegree`, `intermediateGraphDegree`, the graph build algorithm, and its parameters. Use `CUSTOM` only when you have measurements that justify specific values; the defaults derived by cuVS are a better starting point. For background on the parameters themselves, see the [CAGRA indexing guide](/user-guide/api-guides/indexing-guide/cagra) and the [tuning guide](/getting-started/introduction/tuning-indexes). diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index 2b31a361a1..019dda7797 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -101,8 +101,9 @@ public static enum Strategy { * @param writerThreads Number of native cuVS writer threads to use. * @param graphThreads Maximum threads per HNSW graph materialization or serialization operation, * including the calling thread. - * @param graphCopyMemoryBudgetBytes Per-operation ceiling for the raw temporary host adjacency - * copy used by parallel device-graph materialization. + * @param graphCopyMemoryBudgetBytes Configured ceiling for raw temporary host adjacency copies + * used by parallel device-graph materialization. Equal configurations share aggregate + * reservations in this class loader. * @param intermediateGraphDegree The intermediate graph degree while building the CAGRA index. * @param graphdegree The graph degree to use while building the CAGRA index. * @param hnswLayers The number of HNSW layers to build in the HNSW index. From 4b6f8067fa79acad85200359c2749f42fc48d22a Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Wed, 7 Oct 2026 06:24:39 +0000 Subject: [PATCH 19/21] Set graph copy budget default to 24 GiB --- ...vidia-cuvs-lucene-acceleratedhnswparams.md | 42 +++++++++---------- ...nvidia-cuvs-lucene-acceleratedhnswutils.md | 14 +++---- ...om-nvidia-cuvs-lucene-gpubuilthnswgraph.md | 20 ++++----- ...ne-lucene99acceleratedhnswvectorswriter.md | 10 ++--- ...leratedhnswbinaryquantizedvectorswriter.md | 10 ++--- ...leratedhnswscalarquantizedvectorswriter.md | 10 ++--- fern/pages/user_guide/lucene.md | 2 +- .../cuvs/lucene/AcceleratedHNSWParams.java | 5 +-- .../lucene/TestGraphCopyMemoryBudget.java | 17 +++++++- 9 files changed, 72 insertions(+), 58 deletions(-) diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md index b72df83ebe..780646d5a4 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md @@ -292,12 +292,12 @@ public Builder withGraphCopyMemoryBudgetBytes(long graphCopyMemoryBudgetBytes) ``` Set the ceiling for the raw temporary host adjacency copy used by parallel device-graph -materialization. Copies with the same ceiling share aggregate reservations in this class loader. -A differently configured copy cannot overlap active reservations and uses the serial fallback, so -one codec cannot silently raise another codec's active aggregate ceiling. This setting does not -cap the heap-backed Lucene graph, guarantee physical memory availability, or impose a process-wide -limit. A value of `0` preserves the serial fallback without making a temporary device-to-host graph -copy. The default is 42 GiB (`42L << 30` bytes). +materialization. Copies with the same ceiling share aggregate reservations in this class +loader. A copy with a different ceiling cannot overlap those reservations and uses the serial +fallback instead. This setting does not cap the heap-backed Lucene graph, which is allocated +by both the serial and parallel paths, and it is not a guarantee of physical memory +availability or a process-wide limit. A value of `0` disables the temporary copy while +preserving the serial fallback. The default is 24 GiB (`24L << 30` bytes). **Parameters** @@ -309,7 +309,7 @@ copy. The default is 42 GiB (`42L << 30` bytes). instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:403`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:402`_ ### withIntermediateGraphDegree @@ -331,7 +331,7 @@ Default value - \{@value DEFAULT_INT_GRAPH_DEGREE\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:416`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:415`_ ### withGraphDegree @@ -353,7 +353,7 @@ Default value - \{@value DEFAULT_GRAPH_DEGREE\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:429`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:428`_ ### withHNSWLayer @@ -375,7 +375,7 @@ Default value - \{@value DEFAULT_HNSW_LAYERS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:442`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:441`_ ### withMaxConn @@ -397,7 +397,7 @@ Default value - \{@value DEFAULT_MAX_CONN\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:455`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:454`_ ### withBeamWidth @@ -419,7 +419,7 @@ Default value - \{@value DEFAULT_BEAM_WIDTH\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:468`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:467`_ ### withCagraGraphBuildAlgo @@ -440,7 +440,7 @@ Default value - NN_DESCENT instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:480`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:479`_ ### withCuVSIvfPqParams @@ -460,7 +460,7 @@ Set the instance of `CuVSIvfPqParams` instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:491`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:490`_ ### withNumMergeWorkers @@ -481,7 +481,7 @@ Default value - \{@value DEFAULT_NUM_MERGE_WORKERS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:503`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:502`_ ### withMergeExecutorService @@ -502,7 +502,7 @@ Default value an instance with one thread instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:515`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:514`_ ### withStrategy @@ -528,7 +528,7 @@ Default value - HEURISTIC instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:532`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:531`_ ### withCuvsDistanceType @@ -548,7 +548,7 @@ Set the CuvsDistanceType instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:543`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:542`_ ### withNNDescentNumIterations @@ -571,7 +571,7 @@ Default value - \{@value DEFAULT_NN_DESCENT_NUM_ITERATIONS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:557`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:556`_ ### withHnswHeuristicType @@ -595,7 +595,7 @@ the equivalent HNSW graph (graph degree = 2 * maxConn). instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:572`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:571`_ ### build @@ -609,6 +609,6 @@ Create an instance of `AcceleratedHNSWParams` instance of `AcceleratedHNSWParams` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:622`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:621`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:17`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md index 47a7161773..47796c0530 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md @@ -44,7 +44,7 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( int dimensions, CuVSMatrix a Creates a multi-layer HNSW graph from a native matrix without copying the complete dataset to the Java heap. The list view copies only rows selected for an upper layer. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:239`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:271`_ ### writeGraph @@ -71,7 +71,7 @@ a 2D array of offsets | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:416`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:473`_ ### writeMeta @@ -100,7 +100,7 @@ Writes the meta information for the index. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:589`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:646`_ ### printInfoStream @@ -116,7 +116,7 @@ A utility method to print info/debugging messages using InfoStream. | --- | --- | | `msg` | the debugging message to print | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:671`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:728`_ ### writeEmpty @@ -138,7 +138,7 @@ Writes an empty meta information for the field. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:683`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:740`_ ### quantizeFloatVectorsToBinary @@ -161,7 +161,7 @@ Bits are packed: 8 dimensions per byte. A list of byte binary representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:696`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:753`_ ### quantizeFloatVectorsToScalar @@ -181,6 +181,6 @@ Scalar quantization. A list of byte scalar representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:738`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:795`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:35`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md index eb0a751516..f9bd040752 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md @@ -41,7 +41,7 @@ public NodesIterator getNodesOnLevel(int level) Get all nodes on a given level as node 0th ordinals. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:295`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:377`_ ### getNeighbors @@ -62,7 +62,7 @@ Get the neighbors for the node and the level it resides. an instance of NeighborArray -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:313`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:395`_ ### seek @@ -72,7 +72,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Move the pointer to exactly the given level's target. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:338`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:420`_ ### nextNeighbor @@ -82,7 +82,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Iterates over the neighbor list. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:348`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:430`_ ### entryNode @@ -92,7 +92,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns graph's entry point on the top level. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:379`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:461`_ ### maxConn @@ -102,7 +102,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap returns M, the maximum number of connections for a node. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:398`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:480`_ ### neighborCount @@ -112,7 +112,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns the neighbor count. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:413`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:495`_ ### size @@ -122,7 +122,7 @@ public int size() Returns the number of nodes in the graph. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:488`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:570`_ ### numLevels @@ -136,7 +136,7 @@ Returns the number of levels in the HNSW graph. the number of levels -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:497`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:579`_ ### dimensions @@ -150,6 +150,6 @@ Gets the vector dimension. the vector dimension -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:506`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:588`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:27`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md index ea991d5f91..9881ecb1da 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:230`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:231`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:377`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:378`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:386`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:387`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:406`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:407`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:416`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:417`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:55`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md index 2d058c70f6..6875113ad2 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:229`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:230`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:319`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:320`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:350`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:351`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:370`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:371`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:379`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:380`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:57`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md index e4fd6e79ff..bb307322c3 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:256`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:257`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:344`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:345`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:375`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:376`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:395`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:396`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:404`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:405`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:56`_ diff --git a/fern/pages/user_guide/lucene.md b/fern/pages/user_guide/lucene.md index 969ad17ec6..5d1aece839 100644 --- a/fern/pages/user_guide/lucene.md +++ b/fern/pages/user_guide/lucene.md @@ -231,7 +231,7 @@ Both `AcceleratedHNSWParams` and `GPUSearchParams` default to a `HEURISTIC` stra - For the accelerated HNSW codecs, set `maxConn` and `beamWidth`, the HNSW parameters you would tune on the CPU. cuVS derives graph degrees and the build algorithm from them. These two values also configure the CPU fallback writer, so one setting covers both paths. - For the GPU search codec, set `buildQuality`. Higher values spend more build time for a higher-quality graph. This codec also passes `graphDegree` into the heuristic, so leave it at its default unless you intend to cap the graph. -`writerThreads` controls native cuVS build concurrency. The accelerated HNSW codecs default to one writer thread, while the GPU search codec defaults to 32. Accelerated HNSW also offers `graphThreads` for the later CPU graph-materialization and serialization stages; it defaults to one and counts the calling thread. Parallel materialization of a device-backed graph requires a temporary host copy of its INT32 adjacency payload. `graphCopyMemoryBudgetBytes` sets the ceiling for that copy and defaults to 42 GiB (`42L << 30` bytes). Copies with the same ceiling reserve against one shared counter in the codec's class loader. A differently configured copy cannot overlap active reservations and falls back to serial device-row reads, so one codec cannot silently raise another codec's active aggregate ceiling. Configure the same value on every accelerated-HNSW codec when the application requires one consistent classloader-wide policy. +`writerThreads` controls native cuVS build concurrency. The accelerated HNSW codecs default to one writer thread, while the GPU search codec defaults to 32. Accelerated HNSW also offers `graphThreads` for the later CPU graph-materialization and serialization stages; it defaults to one and counts the calling thread. Parallel materialization of a device-backed graph requires a temporary host copy of its INT32 adjacency payload. `graphCopyMemoryBudgetBytes` sets the ceiling for that copy and defaults to 24 GiB (`24L << 30` bytes). Copies with the same ceiling reserve against one shared counter in the codec's class loader. A differently configured copy cannot overlap active reservations and falls back to serial device-row reads, so one codec cannot silently raise another codec's active aggregate ceiling. Configure the same value on every accelerated-HNSW codec when the application requires one consistent classloader-wide policy. The graph-copy budget controls only the extra temporary host copy used by the parallel path. It does not cap the heap-backed Lucene graph (which both paths build), measure current physical memory, guarantee that an admitted native allocation will succeed, or coordinate copies made by separately loaded copies of this library. Long-lived applications such as Solr, Elasticsearch, and OpenSearch should configure every accelerated-HNSW codec consistently and choose a value that leaves headroom for Lucene, the JVM heap, filesystem cache, native libraries, and other components. The core API exposes the setting through `AcceleratedHNSWParams.Builder`; a server integration must map its operator-facing configuration into that builder. Name-only codec SPI construction uses the defaults, including `graphThreads = 1`, for which no temporary graph copy is made. Setting the budget to `0` disables device-graph copying without disabling indexing. diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index 019dda7797..276e821c37 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -53,7 +53,7 @@ public static enum Strategy { public static final int DEFAULT_WRITER_THREADS = 1; public static final int DEFAULT_GRAPH_THREADS = 1; - public static final long DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES = 42L << 30; + public static final long DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES = 24L << 30; public static final int DEFAULT_INT_GRAPH_DEGREE = 128; public static final int DEFAULT_GRAPH_DEGREE = 64; public static final int DEFAULT_HNSW_LAYERS = 1; @@ -395,8 +395,7 @@ public Builder withGraphThreads(int graphThreads) { * fallback instead. This setting does not cap the heap-backed Lucene graph, which is allocated * by both the serial and parallel paths, and it is not a guarantee of physical memory * availability or a process-wide limit. A value of {@code 0} disables the temporary copy while - * preserving the serial fallback. Default value - {@value - * DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES} bytes. + * preserving the serial fallback. The default is 24 GiB ({@code 24L << 30} bytes). * * @param graphCopyMemoryBudgetBytes graph-copy memory budget in bytes * @return instance of {@link Builder} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java index 171673d5c0..ec4b878b91 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java @@ -24,12 +24,22 @@ public class TestGraphCopyMemoryBudget extends LuceneTestCase { public void defaultBudgetAdmitsExpectedBenchmarkShapes() { long defaultBudget = AcceleratedHNSWParams.DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES; long deep100MRequired = GraphCopyMemoryBudget.requiredCopyBytes(100_000_000L, 32); + long deep100MFinalDegree48Required = GraphCopyMemoryBudget.requiredCopyBytes(100_000_000L, 48); + long deep100MDefaultDegreeRequired = GraphCopyMemoryBudget.requiredCopyBytes(100_000_000L, 64); + long deep100MAboveDefaultDegreeRequired = + GraphCopyMemoryBudget.requiredCopyBytes(100_000_000L, 65); long jasper10MRequired = GraphCopyMemoryBudget.requiredCopyBytes(10_000_000L, 32); - assertEquals(42L << 30, defaultBudget); + assertEquals(24L << 30, defaultBudget); assertEquals(12_800_000_000L, deep100MRequired); + assertEquals(19_200_000_000L, deep100MFinalDegree48Required); + assertEquals(25_600_000_000L, deep100MDefaultDegreeRequired); + assertEquals(26_000_000_000L, deep100MAboveDefaultDegreeRequired); assertEquals(1_280_000_000L, jasper10MRequired); assertTrue(deep100MRequired < defaultBudget); + assertTrue(deep100MFinalDegree48Required < defaultBudget); + assertTrue(deep100MDefaultDegreeRequired < defaultBudget); + assertTrue(deep100MAboveDefaultDegreeRequired > defaultBudget); assertTrue(jasper10MRequired < defaultBudget); GraphCopyMemoryBudget budget = new GraphCopyMemoryBudget(); @@ -37,6 +47,11 @@ public void defaultBudgetAdmitsExpectedBenchmarkShapes() { budget.tryReserve(100_000_000L, 32, defaultBudget).orElseThrow()) { // The default admits one 100M-by-32 temporary copy. } + try (GraphCopyMemoryBudget.Reservation ignored = + budget.tryReserve(100_000_000L, 64, defaultBudget).orElseThrow()) { + // The default also admits one 100M graph using the stock degree of 64. + } + assertTrue(budget.tryReserve(100_000_000L, 65, defaultBudget).isEmpty()); } @Test From 26b000ce0d6a248b23cab12d11d074dde211ca45 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Wed, 7 Oct 2026 06:48:20 +0000 Subject: [PATCH 20/21] Document graph copy activation boundary --- fern/pages/user_guide/lucene.md | 2 +- .../lucene/TestCagraIndexParamsFactory.java | 24 +++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/fern/pages/user_guide/lucene.md b/fern/pages/user_guide/lucene.md index 5d1aece839..d99f404de1 100644 --- a/fern/pages/user_guide/lucene.md +++ b/fern/pages/user_guide/lucene.md @@ -231,7 +231,7 @@ Both `AcceleratedHNSWParams` and `GPUSearchParams` default to a `HEURISTIC` stra - For the accelerated HNSW codecs, set `maxConn` and `beamWidth`, the HNSW parameters you would tune on the CPU. cuVS derives graph degrees and the build algorithm from them. These two values also configure the CPU fallback writer, so one setting covers both paths. - For the GPU search codec, set `buildQuality`. Higher values spend more build time for a higher-quality graph. This codec also passes `graphDegree` into the heuristic, so leave it at its default unless you intend to cap the graph. -`writerThreads` controls native cuVS build concurrency. The accelerated HNSW codecs default to one writer thread, while the GPU search codec defaults to 32. Accelerated HNSW also offers `graphThreads` for the later CPU graph-materialization and serialization stages; it defaults to one and counts the calling thread. Parallel materialization of a device-backed graph requires a temporary host copy of its INT32 adjacency payload. `graphCopyMemoryBudgetBytes` sets the ceiling for that copy and defaults to 24 GiB (`24L << 30` bytes). Copies with the same ceiling reserve against one shared counter in the codec's class loader. A differently configured copy cannot overlap active reservations and falls back to serial device-row reads, so one codec cannot silently raise another codec's active aggregate ceiling. Configure the same value on every accelerated-HNSW codec when the application requires one consistent classloader-wide policy. +`writerThreads` controls native cuVS build concurrency. The accelerated HNSW codecs default to one writer thread, while the GPU search codec defaults to 32. Accelerated HNSW also offers `graphThreads` for the later CPU graph-materialization and serialization stages; it defaults to one and counts the calling thread. Parallel materialization is attempted only for graphs with at least 65,536 nodes; smaller graphs use the serial path even when `graphThreads` is greater than one. Parallel materialization of a device-backed graph requires a temporary host copy of its INT32 adjacency payload. `graphCopyMemoryBudgetBytes` sets the ceiling for that copy and defaults to 24 GiB (`24L << 30` bytes). Copies with the same ceiling reserve against one shared counter in the codec's class loader. A differently configured copy cannot overlap active reservations and falls back to serial device-row reads, so one codec cannot silently raise another codec's active aggregate ceiling. Configure the same value on every accelerated-HNSW codec when the application requires one consistent classloader-wide policy. The graph-copy budget controls only the extra temporary host copy used by the parallel path. It does not cap the heap-backed Lucene graph (which both paths build), measure current physical memory, guarantee that an admitted native allocation will succeed, or coordinate copies made by separately loaded copies of this library. Long-lived applications such as Solr, Elasticsearch, and OpenSearch should configure every accelerated-HNSW codec consistently and choose a value that leaves headroom for Lucene, the JVM heap, filesystem cache, native libraries, and other components. The core API exposes the setting through `AcceleratedHNSWParams.Builder`; a server integration must map its operator-facing configuration into that builder. Name-only codec SPI construction uses the defaults, including `graphThreads = 1`, for which no temporary graph copy is made. Setting the budget to `0` disables device-graph copying without disabling indexing. diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java index c4319e6f46..ff37e3f18c 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java @@ -169,6 +169,30 @@ public void testGpuCustomStrategyPassesThroughValues() { assertEquals(12, cagraParams.getNumWriterThreads()); } + /** + * An odd accelerated-HNSW graph degree is a valid CUSTOM input and reaches the copy-budget + * calculation unchanged. Pure Java, no GPU needed. + */ + @Test + public void testHnswCustomOddGraphDegreeReachesCopyBudgetBoundary() { + int graphDegree = 65; + long rows = 100_000_000L; + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withGraphDegree(graphDegree) + .withIntermediateGraphDegree(128) + .build(); + + CagraIndexParams cagraParams = CagraIndexParamsFactory.create(params, rows, 96); + long requiredCopyBytes = + GraphCopyMemoryBudget.requiredCopyBytes(rows, cagraParams.getGraphDegree()); + + assertEquals(graphDegree, cagraParams.getGraphDegree()); + assertEquals(26_000_000_000L, requiredCopyBytes); + assertTrue(requiredCopyBytes > AcceleratedHNSWParams.DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES); + } + /** * The accelerated-HNSW HEURISTIC path delegates to cuVS' native {@code fromHnswParams}, which * derives the graph degrees from maxConn/beamWidth, and re-attaches the caller's writerThreads From cc661705560cbbeb149b02a23a0102e8046c10ed Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Wed, 7 Oct 2026 10:39:26 +0000 Subject: [PATCH 21/21] Finalize graph processing defaults and copy budgets --- ...vidia-cuvs-lucene-acceleratedhnswparams.md | 70 +++++++++-------- fern/pages/user_guide/lucene.md | 4 +- .../cuvs/lucene/AcceleratedHNSWParams.java | 13 +++- .../cuvs/lucene/AcceleratedHNSWUtils.java | 2 +- .../cuvs/lucene/GraphCopyMemoryBudget.java | 27 ++++--- .../lucene/TestAcceleratedHNSWParams.java | 21 ++++- .../lucene/TestGraphCopyMemoryBudget.java | 41 +++++++++- .../TestParallelGraphMaterialization.java | 78 ++++++++++++++++--- .../TestParallelGraphSerialization.java | 32 ++++++-- 9 files changed, 217 insertions(+), 71 deletions(-) diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md index 780646d5a4..c19a81ce31 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md @@ -24,7 +24,7 @@ Get the native cuVS writer threads parameter. cuVS writer threads parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:161`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:163`_ ### getGraphThreads @@ -39,7 +39,7 @@ includes the calling thread; shared helper capacity may reduce actual concurrenc HNSW graph processing threads parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:171`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:173`_ ### getGraphCopyMemoryBudgetBytes @@ -53,7 +53,7 @@ Get the configured temporary host-copy budget for parallel device-graph material graph-copy memory budget in bytes -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:180`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:182`_ ### getIntermediateGraphDegree @@ -67,7 +67,7 @@ Get the intermediate graph degree the graph degree parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:189`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:191`_ ### getGraphdegree @@ -81,7 +81,7 @@ Get the graph degree the graph degree parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:198`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:200`_ ### getHnswLayers @@ -95,7 +95,7 @@ Get the number of HNSW layers the number of HNSW layers -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:207`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:209`_ ### getMaxConn @@ -109,7 +109,7 @@ Get the max connection parameter the max connection parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:216`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:218`_ ### getBeamWidth @@ -123,7 +123,7 @@ Get the beam width parameter the beam width parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:225`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:227`_ ### getCagraGraphBuildAlgo @@ -137,7 +137,7 @@ Get the CAGRA graph build algorithm the CAGRA graph build algorithm -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:234`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:236`_ ### getCuVSIvfPqParams @@ -151,7 +151,7 @@ Get the instance of `CuVSIvfPqParams` the instance of `CuVSIvfPqParams` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:243`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:245`_ ### getNumMergeWorkers @@ -165,7 +165,7 @@ Get the number of merge workers set to be used in the fallback mechanism the number of merge workers -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:252`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:254`_ ### getMergeExec @@ -179,7 +179,7 @@ Get the instance of the `ExecutorService` to be used in the fallback mechanism the instance of the `ExecutorService` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:261`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:263`_ ### getStrategy @@ -196,7 +196,7 @@ When CUSTOM is chosen, the build algorithm and its parameters (either defaults o get the chosen `Strategy` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:273`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:275`_ ### getCuvsDistanceType @@ -210,7 +210,7 @@ Get the cuvs distance type the distance type -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:282`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:284`_ ### getNNDescentNumIterations @@ -224,7 +224,7 @@ get the number of Iterations to run if building with NN_DESCENT the number of iterations for NN_DESCENT -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:291`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:293`_ ### getHnswHeuristicType @@ -239,7 +239,7 @@ beamWidth. Only consulted under the `Strategy#HEURISTIC` strategy. the `HnswHeuristicType` to hand to cuVS -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:301`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:303`_ ### withWriterThreads @@ -261,7 +261,7 @@ Default value - \{@value DEFAULT_WRITER_THREADS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:372`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:374`_ ### withGraphThreads @@ -283,7 +283,7 @@ includes the calling thread. Valid range - Minimum: \{@value MIN_GRAPH_THREADS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:385`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:387`_ ### withGraphCopyMemoryBudgetBytes @@ -297,7 +297,9 @@ loader. A copy with a different ceiling cannot overlap those reservations and us fallback instead. This setting does not cap the heap-backed Lucene graph, which is allocated by both the serial and parallel paths, and it is not a guarantee of physical memory availability or a process-wide limit. A value of `0` disables the temporary copy while -preserving the serial fallback. The default is 24 GiB (`24L << 30` bytes). +preserving the serial fallback. A value of `-1` removes the copy ceiling; reservations +remain accounted, and an unlimited policy cannot overlap a different active policy. Values +less than `-1` are rejected. The default is 24 GiB (`24L << 30` bytes). **Parameters** @@ -309,7 +311,7 @@ preserving the serial fallback. The default is 24 GiB (`24L << 30` bytes). instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:402`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:406`_ ### withIntermediateGraphDegree @@ -331,7 +333,7 @@ Default value - \{@value DEFAULT_INT_GRAPH_DEGREE\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:415`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:419`_ ### withGraphDegree @@ -353,7 +355,7 @@ Default value - \{@value DEFAULT_GRAPH_DEGREE\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:428`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:432`_ ### withHNSWLayer @@ -375,7 +377,7 @@ Default value - \{@value DEFAULT_HNSW_LAYERS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:441`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:445`_ ### withMaxConn @@ -397,7 +399,7 @@ Default value - \{@value DEFAULT_MAX_CONN\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:454`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:458`_ ### withBeamWidth @@ -419,7 +421,7 @@ Default value - \{@value DEFAULT_BEAM_WIDTH\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:467`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:471`_ ### withCagraGraphBuildAlgo @@ -440,7 +442,7 @@ Default value - NN_DESCENT instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:479`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:483`_ ### withCuVSIvfPqParams @@ -460,7 +462,7 @@ Set the instance of `CuVSIvfPqParams` instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:490`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:494`_ ### withNumMergeWorkers @@ -481,7 +483,7 @@ Default value - \{@value DEFAULT_NUM_MERGE_WORKERS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:502`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:506`_ ### withMergeExecutorService @@ -502,7 +504,7 @@ Default value an instance with one thread instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:514`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:518`_ ### withStrategy @@ -528,7 +530,7 @@ Default value - HEURISTIC instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:531`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:535`_ ### withCuvsDistanceType @@ -548,7 +550,7 @@ Set the CuvsDistanceType instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:542`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:546`_ ### withNNDescentNumIterations @@ -571,7 +573,7 @@ Default value - \{@value DEFAULT_NN_DESCENT_NUM_ITERATIONS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:556`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:560`_ ### withHnswHeuristicType @@ -595,7 +597,7 @@ the equivalent HNSW graph (graph degree = 2 * maxConn). instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:571`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:575`_ ### build @@ -609,6 +611,6 @@ Create an instance of `AcceleratedHNSWParams` instance of `AcceleratedHNSWParams` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:621`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:626`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:17`_ diff --git a/fern/pages/user_guide/lucene.md b/fern/pages/user_guide/lucene.md index d99f404de1..583e15ded1 100644 --- a/fern/pages/user_guide/lucene.md +++ b/fern/pages/user_guide/lucene.md @@ -231,9 +231,9 @@ Both `AcceleratedHNSWParams` and `GPUSearchParams` default to a `HEURISTIC` stra - For the accelerated HNSW codecs, set `maxConn` and `beamWidth`, the HNSW parameters you would tune on the CPU. cuVS derives graph degrees and the build algorithm from them. These two values also configure the CPU fallback writer, so one setting covers both paths. - For the GPU search codec, set `buildQuality`. Higher values spend more build time for a higher-quality graph. This codec also passes `graphDegree` into the heuristic, so leave it at its default unless you intend to cap the graph. -`writerThreads` controls native cuVS build concurrency. The accelerated HNSW codecs default to one writer thread, while the GPU search codec defaults to 32. Accelerated HNSW also offers `graphThreads` for the later CPU graph-materialization and serialization stages; it defaults to one and counts the calling thread. Parallel materialization is attempted only for graphs with at least 65,536 nodes; smaller graphs use the serial path even when `graphThreads` is greater than one. Parallel materialization of a device-backed graph requires a temporary host copy of its INT32 adjacency payload. `graphCopyMemoryBudgetBytes` sets the ceiling for that copy and defaults to 24 GiB (`24L << 30` bytes). Copies with the same ceiling reserve against one shared counter in the codec's class loader. A differently configured copy cannot overlap active reservations and falls back to serial device-row reads, so one codec cannot silently raise another codec's active aggregate ceiling. Configure the same value on every accelerated-HNSW codec when the application requires one consistent classloader-wide policy. +`writerThreads` controls native cuVS build concurrency. The accelerated HNSW codecs default to one writer thread, while the GPU search codec defaults to 32. Accelerated HNSW also offers `graphThreads` for the later CPU graph-materialization and serialization stages; it defaults to 16 and counts the calling thread. Parallel materialization is attempted only for graphs with at least 65,536 nodes; smaller graphs use the serial path even when `graphThreads` is greater than one. Parallel materialization of a device-backed graph requires a temporary host copy of its INT32 adjacency payload. `graphCopyMemoryBudgetBytes` sets the ceiling for that copy and defaults to 24 GiB (`24L << 30` bytes). Copies with the same ceiling reserve against one shared counter in the codec's class loader. A differently configured copy cannot overlap active reservations and falls back to serial device-row reads, so one codec cannot silently raise another codec's active aggregate ceiling. Configure the same value on every accelerated-HNSW codec when the application requires one consistent classloader-wide policy. -The graph-copy budget controls only the extra temporary host copy used by the parallel path. It does not cap the heap-backed Lucene graph (which both paths build), measure current physical memory, guarantee that an admitted native allocation will succeed, or coordinate copies made by separately loaded copies of this library. Long-lived applications such as Solr, Elasticsearch, and OpenSearch should configure every accelerated-HNSW codec consistently and choose a value that leaves headroom for Lucene, the JVM heap, filesystem cache, native libraries, and other components. The core API exposes the setting through `AcceleratedHNSWParams.Builder`; a server integration must map its operator-facing configuration into that builder. Name-only codec SPI construction uses the defaults, including `graphThreads = 1`, for which no temporary graph copy is made. Setting the budget to `0` disables device-graph copying without disabling indexing. +The graph-copy budget controls only the extra temporary host copy used by the parallel path. It does not cap the heap-backed Lucene graph (which both paths build), measure current physical memory, guarantee that an admitted native allocation will succeed, or coordinate copies made by separately loaded copies of this library. Long-lived applications such as Solr, Elasticsearch, and OpenSearch should configure every accelerated-HNSW codec consistently and choose a value that leaves headroom for Lucene, the JVM heap, filesystem cache, native libraries, and other components. The core API exposes the setting through `AcceleratedHNSWParams.Builder`; a server integration must map its operator-facing configuration into that builder. Name-only codec SPI construction uses the defaults, including `graphThreads = 16`, so eligible device-backed graphs attempt the temporary copy unless the 24-GiB budget denies it. Setting the budget to `0` disables device-graph copying without disabling indexing. Setting it to `-1` removes the ceiling, which allows every eligible concurrent copy to be attempted and transfers the resulting native-allocation, memory-pressure, and OOM risk to the application. Values less than `-1` are rejected. Switching either class to the `CUSTOM` strategy exposes the underlying CAGRA parameters directly, including `graphDegree`, `intermediateGraphDegree`, the graph build algorithm, and its parameters. Use `CUSTOM` only when you have measurements that justify specific values; the defaults derived by cuVS are a better starting point. For background on the parameters themselves, see the [CAGRA indexing guide](/user-guide/api-guides/indexing-guide/cagra) and the [tuning guide](/getting-started/introduction/tuning-indexes). diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index 276e821c37..0e9ae7d024 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -52,7 +52,9 @@ public static enum Strategy { public static final int MAX_NN_DESCENT_NUM_ITERATIONS = 100; public static final int DEFAULT_WRITER_THREADS = 1; - public static final int DEFAULT_GRAPH_THREADS = 1; + public static final int DEFAULT_GRAPH_THREADS = 16; + public static final long UNLIMITED_GRAPH_COPY_MEMORY_BUDGET_BYTES = -1L; + public static final long DISABLED_GRAPH_COPY_MEMORY_BUDGET_BYTES = 0L; public static final long DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES = 24L << 30; public static final int DEFAULT_INT_GRAPH_DEGREE = 128; public static final int DEFAULT_GRAPH_DEGREE = 64; @@ -395,7 +397,9 @@ public Builder withGraphThreads(int graphThreads) { * fallback instead. This setting does not cap the heap-backed Lucene graph, which is allocated * by both the serial and parallel paths, and it is not a guarantee of physical memory * availability or a process-wide limit. A value of {@code 0} disables the temporary copy while - * preserving the serial fallback. The default is 24 GiB ({@code 24L << 30} bytes). + * preserving the serial fallback. A value of {@code -1} removes the copy ceiling; reservations + * remain accounted, and an unlimited policy cannot overlap a different active policy. Values + * less than {@code -1} are rejected. The default is 24 GiB ({@code 24L << 30} bytes). * * @param graphCopyMemoryBudgetBytes graph-copy memory budget in bytes * @return instance of {@link Builder} @@ -584,8 +588,9 @@ private void validate() throws IllegalArgumentException { "writerThreads", writerThreads, MIN_WRITER_THREADS, MAX_WRITER_THREADS); ParameterValidation.checkRange( "graphThreads", graphThreads, MIN_GRAPH_THREADS, MAX_GRAPH_THREADS); - if (graphCopyMemoryBudgetBytes < 0) { - throw new IllegalArgumentException("graphCopyMemoryBudgetBytes must be non-negative."); + if (graphCopyMemoryBudgetBytes < UNLIMITED_GRAPH_COPY_MEMORY_BUDGET_BYTES) { + throw new IllegalArgumentException( + "graphCopyMemoryBudgetBytes must be -1 (unlimited) or non-negative."); } ParameterValidation.checkRange( "intermediateGraphDegree", intermediateGraphDegree, MIN_INT_GRAPH_DEG, MAX_INT_GRAPH_DEG); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index 3d27e7f717..c7b9e1836b 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -473,7 +473,7 @@ private static CuVSMatrix buildCagraGraphForSubset( */ public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) throws IOException { - return writeGraph(graph, vectorIndex, AcceleratedHNSWParams.DEFAULT_GRAPH_THREADS); + return writeGraph(graph, vectorIndex, /* graphThreads= */ 1); } static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int graphThreads) diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java index 9959fd3781..cba990fc27 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java @@ -8,10 +8,11 @@ /** Coordinates temporary native graph copies across concurrent segment flushes. */ final class GraphCopyMemoryBudget { + private static final long NO_ACTIVE_BUDGET_BYTES = Long.MIN_VALUE; private static final GraphCopyMemoryBudget SHARED = new GraphCopyMemoryBudget(); private long reservedCopyBytes; - private long activeBudgetBytes = -1; + private long activeBudgetBytes = NO_ACTIVE_BUDGET_BYTES; static GraphCopyMemoryBudget shared() { return SHARED; @@ -20,24 +21,32 @@ static GraphCopyMemoryBudget shared() { /** * Tries to reserve the raw INT32 payload of one temporary device-to-host adjacency copy. * Reservations are shared by callers in this class loader. Overlapping reservations must use - * the same configured ceiling; a caller with a different ceiling is denied until the active - * reservations are released. This keeps one caller from silently raising another caller's - * active aggregate ceiling. + * the same configured budget; a caller with a different budget is denied until the active + * reservations are released. A budget of {@code -1} is unlimited, but its reservations remain + * accounted and cannot overlap a finite policy. This keeps one caller from silently replacing + * another caller's active policy. */ synchronized Optional tryReserve( long rows, long columns, long configuredBudgetBytes) { long requiredCopyBytes = requiredCopyBytes(rows, columns); - if (requiredCopyBytes < 0 || configuredBudgetBytes < 0) { + if (requiredCopyBytes < 0 + || configuredBudgetBytes < AcceleratedHNSWParams.UNLIMITED_GRAPH_COPY_MEMORY_BUDGET_BYTES) { return Optional.empty(); } - if (requiredCopyBytes > configuredBudgetBytes) { + boolean unlimited = + configuredBudgetBytes == AcceleratedHNSWParams.UNLIMITED_GRAPH_COPY_MEMORY_BUDGET_BYTES; + if (!unlimited && requiredCopyBytes > configuredBudgetBytes) { return Optional.empty(); } if (reservedCopyBytes != 0 && activeBudgetBytes != configuredBudgetBytes) { return Optional.empty(); } - if (reservedCopyBytes > configuredBudgetBytes - || requiredCopyBytes > configuredBudgetBytes - reservedCopyBytes) { + if (requiredCopyBytes > Long.MAX_VALUE - reservedCopyBytes) { + return Optional.empty(); + } + if (!unlimited + && (reservedCopyBytes > configuredBudgetBytes + || requiredCopyBytes > configuredBudgetBytes - reservedCopyBytes)) { return Optional.empty(); } if (reservedCopyBytes == 0) { @@ -65,7 +74,7 @@ private synchronized void release(Reservation reservation) { } reservedCopyBytes -= reservation.copyBytes; if (reservedCopyBytes == 0) { - activeBudgetBytes = -1; + activeBudgetBytes = NO_ACTIVE_BUDGET_BYTES; } reservation.released = true; } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java index 748b13c555..be127878d4 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java @@ -18,6 +18,7 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_NUM_MERGE_WORKERS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_STRATEGY; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_WRITER_THREADS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DISABLED_GRAPH_COPY_MEMORY_BUDGET_BYTES; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_BEAM_WIDTH; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_GRAPH_DEG; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_GRAPH_THREADS; @@ -36,6 +37,7 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_NN_DESCENT_NUM_ITERATIONS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_NUM_MERGE_WORKERS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_WRITER_THREADS; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.UNLIMITED_GRAPH_COPY_MEMORY_BUDGET_BYTES; import static java.lang.Integer.MAX_VALUE; import static java.lang.Integer.MIN_VALUE; @@ -60,6 +62,7 @@ public void testAcceleratedHNSWParamsDefaultValues() { assertEquals(DEFAULT_BEAM_WIDTH, params.getBeamWidth()); assertEquals(DEFAULT_GRAPH_DEGREE, params.getGraphdegree()); assertEquals(DEFAULT_GRAPH_COPY_MEMORY_BUDGET_BYTES, params.getGraphCopyMemoryBudgetBytes()); + assertEquals(16, DEFAULT_GRAPH_THREADS); assertEquals(DEFAULT_GRAPH_THREADS, params.getGraphThreads()); assertEquals(DEFAULT_HNSW_LAYERS, params.getHnswLayers()); assertEquals(DEFAULT_INT_GRAPH_DEGREE, params.getIntermediateGraphDegree()); @@ -230,9 +233,15 @@ public void testAcceleratedHNSWParamsInvalidGraphThreads() { @Test public void testAcceleratedHNSWParamsGraphCopyMemoryBudget() { assertEquals( - 0L, + UNLIMITED_GRAPH_COPY_MEMORY_BUDGET_BYTES, new AcceleratedHNSWParams.Builder() - .withGraphCopyMemoryBudgetBytes(0) + .withGraphCopyMemoryBudgetBytes(UNLIMITED_GRAPH_COPY_MEMORY_BUDGET_BYTES) + .build() + .getGraphCopyMemoryBudgetBytes()); + assertEquals( + DISABLED_GRAPH_COPY_MEMORY_BUDGET_BYTES, + new AcceleratedHNSWParams.Builder() + .withGraphCopyMemoryBudgetBytes(DISABLED_GRAPH_COPY_MEMORY_BUDGET_BYTES) .build() .getGraphCopyMemoryBudgetBytes()); assertEquals( @@ -243,7 +252,13 @@ public void testAcceleratedHNSWParamsGraphCopyMemoryBudget() { .getGraphCopyMemoryBudgetBytes()); assertThrows( IllegalArgumentException.class, - () -> new AcceleratedHNSWParams.Builder().withGraphCopyMemoryBudgetBytes(-1).build()); + () -> new AcceleratedHNSWParams.Builder().withGraphCopyMemoryBudgetBytes(-2).build()); + assertThrows( + IllegalArgumentException.class, + () -> + new AcceleratedHNSWParams.Builder() + .withGraphCopyMemoryBudgetBytes(Long.MIN_VALUE) + .build()); } @Test diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java index ec4b878b91..5035569ee3 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java @@ -131,11 +131,47 @@ public void overlappingReservationsCannotMixConfiguredCeilings() { try (GraphCopyMemoryBudget.Reservation ignored = budget.tryReserve(300, 1, largeCeiling).orElseThrow()) { assertTrue(budget.tryReserve(100, 1, smallCeiling).isEmpty()); + assertTrue( + budget + .tryReserve(100, 1, AcceleratedHNSWParams.UNLIMITED_GRAPH_COPY_MEMORY_BUDGET_BYTES) + .isEmpty()); + } + + try (GraphCopyMemoryBudget.Reservation first = + budget + .tryReserve(100, 1, AcceleratedHNSWParams.UNLIMITED_GRAPH_COPY_MEMORY_BUDGET_BYTES) + .orElseThrow(); + GraphCopyMemoryBudget.Reservation second = + budget + .tryReserve(300, 1, AcceleratedHNSWParams.UNLIMITED_GRAPH_COPY_MEMORY_BUDGET_BYTES) + .orElseThrow()) { + assertTrue(budget.tryReserve(100, 1, smallCeiling).isEmpty()); } try (GraphCopyMemoryBudget.Reservation ignored = budget.tryReserve(100, 1, smallCeiling).orElseThrow()) { - // Releasing all reservations resets the active ceiling. + // Releasing all unlimited reservations resets the active policy. + } + } + + @Test + public void unlimitedReservationsFailClosedOnAccountingOverflow() { + long rows = Integer.MAX_VALUE; + long columns = 1_000_000_000L; + long requiredCopyBytes = GraphCopyMemoryBudget.requiredCopyBytes(rows, columns); + assertTrue(requiredCopyBytes > Long.MAX_VALUE / 2); + + GraphCopyMemoryBudget budget = new GraphCopyMemoryBudget(); + try (GraphCopyMemoryBudget.Reservation ignored = + budget + .tryReserve( + rows, columns, AcceleratedHNSWParams.UNLIMITED_GRAPH_COPY_MEMORY_BUDGET_BYTES) + .orElseThrow()) { + assertTrue( + budget + .tryReserve( + rows, columns, AcceleratedHNSWParams.UNLIMITED_GRAPH_COPY_MEMORY_BUDGET_BYTES) + .isEmpty()); } } @@ -172,6 +208,7 @@ public void invalidShapesAndBudgetsFailClosed() { assertTrue(budget.tryReserve(1, 0, Long.MAX_VALUE).isEmpty()); assertTrue(budget.tryReserve(-1, 1, Long.MAX_VALUE).isEmpty()); assertTrue(budget.tryReserve(1, -1, Long.MAX_VALUE).isEmpty()); - assertTrue(budget.tryReserve(1, 1, -1).isEmpty()); + assertTrue(budget.tryReserve(1, 1, -2).isEmpty()); + assertTrue(budget.tryReserve(1, 1, Long.MIN_VALUE).isEmpty()); } } diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java index d263c9896a..13860fdbd1 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java @@ -14,6 +14,8 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicInteger; import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.util.hnsw.HnswGraph; @@ -49,10 +51,21 @@ rows, new AtomicInteger(), null, executionProbe)) { @Test public void overflowingDeviceShapeUsesSerialFallback() throws Exception { int[][] adjacency = IntGraphTestMatrix.randomRows(NUM_NODES, 1, 0); - try (CuVSMatrix matrix = new IntGraphTestMatrix.DeviceMatrix(adjacency, Long.MAX_VALUE)) { + Set sourceReadThreads = ConcurrentHashMap.newKeySet(); + try (CuVSMatrix matrix = + new IntGraphTestMatrix.DeviceMatrix(adjacency, Long.MAX_VALUE) { + @Override + public RowView getRow(long row) { + sourceReadThreads.add(Thread.currentThread()); + return super.getRow(row); + } + }) { GPUBuiltHnswGraph graph = newSingleLayerGraph(matrix, GRAPH_THREADS); - assertEquals(NUM_NODES, graph.size()); + for (int node = 0; node < NUM_NODES; node++) { + assertArrayEquals(adjacency[node], arcsOf(graph, 0, node)); + } } + assertEquals(1, sourceReadThreads.size()); } @Test @@ -211,8 +224,15 @@ public void toHost(CuVSHostMatrix target) { public void configuredBudgetControlsWhetherDeviceCopyRuns() throws Exception { int[][] sourceRows = IntGraphTestMatrix.randomRows(NUM_NODES, DEGREE, 5); AtomicInteger copyCount = new AtomicInteger(); + Set sourceReadThreads = ConcurrentHashMap.newKeySet(); CuVSDeviceMatrix source = new IntGraphTestMatrix.DeviceMatrix(sourceRows, DEGREE) { + @Override + public RowView getRow(long row) { + sourceReadThreads.add(Thread.currentThread()); + return super.getRow(row); + } + @Override public void toHost(CuVSHostMatrix target) { copyCount.incrementAndGet(); @@ -221,20 +241,60 @@ public void toHost(CuVSHostMatrix target) { long requiredCopyBytes = GraphCopyMemoryBudget.requiredCopyBytes(NUM_NODES, DEGREE); GraphCopyMemoryBudget budget = new GraphCopyMemoryBudget(); - GPUBuiltHnswGraph.materializeDeviceAdjacency( - source, - NUM_NODES, - GRAPH_THREADS, - budget, - requiredCopyBytes - 1, - () -> new IntGraphTestMatrix.TrackingHostMatrix(new AtomicInteger(), null)); + NeighborArray[] disabledResult = + GPUBuiltHnswGraph.materializeDeviceAdjacency( + source, + NUM_NODES, + GRAPH_THREADS, + budget, + AcceleratedHNSWParams.DISABLED_GRAPH_COPY_MEMORY_BUDGET_BYTES, + TestParallelGraphMaterialization::failUnexpectedHostCopyAllocation); + assertAdjacencyEquals(sourceRows, disabledResult); + assertEquals(0, copyCount.get()); + assertEquals(1, sourceReadThreads.size()); + + sourceReadThreads.clear(); + NeighborArray[] insufficientResult = + GPUBuiltHnswGraph.materializeDeviceAdjacency( + source, + NUM_NODES, + GRAPH_THREADS, + budget, + requiredCopyBytes - 1, + TestParallelGraphMaterialization::failUnexpectedHostCopyAllocation); + assertAdjacencyEquals(sourceRows, insufficientResult); assertEquals(0, copyCount.get()); + assertEquals(1, sourceReadThreads.size()); CuVSHostMatrix hostCopy = new IntGraphTestMatrix.TrackingHostMatrix(sourceRows, new AtomicInteger(), null, null); GPUBuiltHnswGraph.materializeDeviceAdjacency( source, NUM_NODES, GRAPH_THREADS, budget, requiredCopyBytes, () -> hostCopy); assertEquals(1, copyCount.get()); + + CuVSHostMatrix unlimitedHostCopy = + new IntGraphTestMatrix.TrackingHostMatrix(sourceRows, new AtomicInteger(), null, null); + GPUBuiltHnswGraph.materializeDeviceAdjacency( + source, + NUM_NODES, + GRAPH_THREADS, + budget, + AcceleratedHNSWParams.UNLIMITED_GRAPH_COPY_MEMORY_BUDGET_BYTES, + () -> unlimitedHostCopy); + assertEquals(2, copyCount.get()); + } + + private static CuVSHostMatrix failUnexpectedHostCopyAllocation() { + throw new AssertionError("denied device adjacency must not allocate a host copy"); + } + + private static void assertAdjacencyEquals(int[][] expectedRows, NeighborArray[] actualNeighbors) { + assertEquals(expectedRows.length, actualNeighbors.length); + for (int node = 0; node < expectedRows.length; node++) { + assertArrayEquals( + expectedRows[node], + Arrays.copyOf(actualNeighbors[node].nodes(), actualNeighbors[node].size())); + } } private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, int numThreads) diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java index ac6c1fecc4..bec72e0ff1 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java @@ -8,6 +8,8 @@ import java.io.IOException; import java.util.Arrays; import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import org.apache.lucene.store.ByteBuffersDirectory; import org.apache.lucene.store.Directory; import org.apache.lucene.store.IOContext; @@ -28,11 +30,12 @@ public class TestParallelGraphSerialization extends LuceneTestCase { public void parallelSerializationMatchesSerial() throws Exception { try (CuVSMatrix matrix = IntGraphTestMatrix.random(NUM_NODES, DEGREE, 2); Directory dir = new ByteBuffersDirectory()) { - GPUBuiltHnswGraph serialGraph = newSingleLayerGraph(matrix); + SerialRecordingGraph serialGraph = new SerialRecordingGraph(matrix); IntGraphTestMatrix.ParallelExecutionProbe executionProbe = new IntGraphTestMatrix.ParallelExecutionProbe(); GPUBuiltHnswGraph parallelGraph = new RecordingGraph(matrix, executionProbe); assertSerialAndParallelMatch(serialGraph, parallelGraph, dir); + assertEquals(1, serialGraph.threadCount()); assertTrue(executionProbe.threadCount() > 1); } } @@ -99,7 +102,9 @@ private static void assertSerialAndParallelMatch( } int[][] parallelOffsets; try (IndexOutput out = dir.createOutput("parallel", IOContext.DEFAULT)) { - parallelOffsets = AcceleratedHNSWUtils.writeGraph(parallelGraph, out, GRAPH_THREADS); + parallelOffsets = + AcceleratedHNSWUtils.writeGraph( + parallelGraph, out, AcceleratedHNSWParams.DEFAULT_GRAPH_THREADS); } assertEquals(serialOffsets.length, parallelOffsets.length); @@ -109,11 +114,6 @@ private static void assertSerialAndParallelMatch( assertArrayEquals(readAllBytes(dir, "serial"), readAllBytes(dir, "parallel")); } - private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency) { - return new GPUBuiltHnswGraph( - NUM_NODES, /* dimensions= */ 4, Arrays.asList((int[]) null), List.of(layer0Adjacency)); - } - private static byte[] readAllBytes(Directory dir, String name) throws Exception { try (IndexInput in = dir.openInput(name, IOContext.DEFAULT)) { byte[] bytes = new byte[(int) in.length()]; @@ -138,6 +138,24 @@ public NeighborArray getNeighbors(int level, int node) { } } + private static final class SerialRecordingGraph extends GPUBuiltHnswGraph { + private final Set threads = ConcurrentHashMap.newKeySet(); + + SerialRecordingGraph(CuVSMatrix layer0Adjacency) { + super(NUM_NODES, /* dimensions= */ 4, Arrays.asList((int[]) null), List.of(layer0Adjacency)); + } + + @Override + public NeighborArray getNeighbors(int level, int node) { + threads.add(Thread.currentThread()); + return super.getNeighbors(level, node); + } + + int threadCount() { + return threads.size(); + } + } + /** * Crosses a maxConn-derived byte-bounded wave with sparse rows and no retained heap graph. */