From 399426c0372ef7aa1c4a1075c19b2d89d8beb1a5 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Tue, 28 Apr 2026 19:27:43 -0400 Subject: [PATCH 1/8] Initial work --- .../cuvs/lucene/AcceleratedHNSWUtils.java | 25 ++- .../cuvs/lucene/CuvsResourcesManager.java | 202 ++++++++++++++++++ .../cuvs/lucene/FilterCuVSProvider.java | 168 --------------- .../lucene/FilterCuVSServiceProvider.java | 24 --- .../Lucene99AcceleratedHNSWVectorsFormat.java | 9 +- .../Lucene99AcceleratedHNSWVectorsWriter.java | 55 +++-- ...ratedHNSWBinaryQuantizedVectorsFormat.java | 8 +- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 61 ++++-- ...ratedHNSWScalarQuantizedVectorsFormat.java | 8 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 61 ++++-- 10 files changed, 357 insertions(+), 264 deletions(-) create mode 100644 src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java delete mode 100644 src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java delete mode 100644 src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java diff --git a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index ac4c9426..338402a3 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -13,6 +13,7 @@ import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.CuVSIvfPqParams; import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.RowView; import java.io.IOException; import java.util.ArrayList; @@ -89,7 +90,8 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( int hnswLayers, int graphDegree, CagraIndexParams params, - QuantizationType quantization) + QuantizationType quantization, + CuVSResources cuVSResources) throws Throwable { // Calculate M as cagraGraphDegree/2 @@ -143,7 +145,13 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( // Build CAGRA graph for this layer layerAdjacencies.add( buildCagraGraphForSubset( - selectedVectors, selectedNodes, 0, params, dimensions, quantization)); + selectedVectors, + selectedNodes, + 0, + params, + dimensions, + quantization, + cuVSResources)); } else { @@ -157,7 +165,13 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( // Build CAGRA graph for this layer layerAdjacencies.add( buildCagraGraphForSubset( - selectedVectors, selectedNodes, bytesPerVector, params, dimensions, quantization)); + selectedVectors, + selectedNodes, + bytesPerVector, + params, + dimensions, + quantization, + cuVSResources)); } // Update for next iteration @@ -181,7 +195,8 @@ private static CuVSMatrix buildCagraGraphForSubset( int bytesPerVector, CagraIndexParams params, int dimensions, - QuantizationType quantization) + QuantizationType quantization, + CuVSResources cuVSResources) throws Throwable { CuVSMatrix subsetDataset; @@ -198,7 +213,7 @@ private static CuVSMatrix buildCagraGraphForSubset( // Build CAGRA index for the subset CagraIndex subsetIndex = - CagraIndex.newBuilder(getCuVSResourcesInstance()) + CagraIndex.newBuilder(cuVSResources) .withDataset(subsetDataset) .withIndexParams(params) .build(); diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java b/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java new file mode 100644 index 00000000..fc80e8cd --- /dev/null +++ b/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java @@ -0,0 +1,202 @@ +/* + * 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.CagraIndexParams.CagraGraphBuildAlgo.IVF_PQ; +import static com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT; + +import com.nvidia.cuvs.CagraIndexParams; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CuVSIvfPqIndexParams; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.GPUInfoProvider; +import com.nvidia.cuvs.spi.CuVSProvider; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.locks.Condition; +import java.util.concurrent.locks.ReentrantLock; +import java.util.logging.Level; +import java.util.logging.Logger; + +public class CuvsResourcesManager { + + private static final Logger log = Logger.getLogger(Utils.class.getName()); + private static final CuVSProvider PROVIDER = CuVSProvider.provider(); + private static final GPUInfoProvider GPU_INFO_PROVIDER = PROVIDER.gpuInfoProvider(); + private static final int MAX_POOL_SIZE = 512; + + private ManagedCuVSResources[] pool; + private ReentrantLock lock; + private Condition resourcesAvailable; + private AtomicLong reserveMemory; + private long totalDeviceMemory; + private int capacity; + + public CuvsResourcesManager(int capacity) { + if (capacity > MAX_POOL_SIZE || capacity <= 0) { + throw new IllegalArgumentException( + "Invalid capacity, should be between 1 and " + MAX_POOL_SIZE); + } + this.capacity = capacity; + pool = new ManagedCuVSResources[capacity]; + lock = new ReentrantLock(); + resourcesAvailable = lock.newCondition(); + for (int i = 0; i < capacity; i++) { + pool[i] = new ManagedCuVSResources(getCuVSResourceInstance()); + } + CuVSResources rx = getCuVSResourceInstance(); + reserveMemory = new AtomicLong(); + totalDeviceMemory = GPU_INFO_PROVIDER.getCurrentInfo(rx).totalDeviceMemoryInBytes(); + rx.close(); + } + + private ManagedCuVSResources getAvailableResourcesFromPool() { + try { + lock.lock(); + for (int i = 0; i < capacity; i++) { + if (pool[i] != null && !pool[i].isLocked()) { + return pool[i]; + } + } + } finally { + lock.unlock(); + } + return null; + } + + private int getNumLockedResources() { + try { + lock.lock(); + int res = 0; + for (int i = 0; i < capacity; i++) { + if (pool[i].isLocked()) { + res += 1; + } + } + return res; + } finally { + lock.unlock(); + } + } + + public ManagedCuVSResources acquireResource(long rows, long dimension, CagraIndexParams params) + throws InterruptedException { + try { + lock.lock(); + long neededMem = getEstimatedMemoryRequirement(rows, dimension, params); + + if (neededMem > totalDeviceMemory) { + throw new RuntimeException("Not enough GPU device memory available"); + } + + while (getNumLockedResources() == capacity + || (totalDeviceMemory - reserveMemory.get()) < neededMem) { + resourcesAvailable.await(); + } + reserveMemory.addAndGet(neededMem); + + ManagedCuVSResources res = getAvailableResourcesFromPool(); + assert res != null; + + res.setNeededMemory(neededMem); + res.lock(); + return res; + } finally { + lock.unlock(); + } + } + + public void releaseResource(ManagedCuVSResources resource) { + try { + lock.lock(); + reserveMemory.addAndGet(-resource.getNeededMemory()); + resource.unlock(); + resourcesAvailable.signalAll(); + } finally { + lock.unlock(); + } + } + + public void shutdown() { + try { + lock.lock(); + for (int i = 0; i < capacity; i++) { + if (pool[i] != null && !pool[i].isLocked() && pool[i].getResource() != null) { + pool[i].getResource().close(); + } + } + } finally { + lock.unlock(); + } + } + + private static CuVSResources getCuVSResourceInstance() { + try { + return CuVSResources.create(); + } catch (UnsupportedOperationException uoe) { + log.log( + Level.WARNING, + "cuVS is not supported on this platform or java version: " + uoe.getMessage()); + } catch (Throwable t) { + if (t instanceof ExceptionInInitializerError ex) { + t = ex.getCause(); + } + log.log(Level.WARNING, "Exception occurred during creation of cuVS resources. " + t); + } + return null; + } + + private long getEstimatedMemoryRequirement(long rows, long dimension, CagraIndexParams params) { + CagraGraphBuildAlgo buildAlgo = params.getCagraGraphBuildAlgo(); + if (buildAlgo.equals(NN_DESCENT)) { + return 2 * rows * dimension * Float.BYTES; + } else if (buildAlgo.equals(IVF_PQ)) { + CuVSIvfPqIndexParams ip = params.getCuVSIvfPqParams().getIndexParams(); + long approximatedIvfBytes = + (long) + (rows * (ip.getPqDim() * (ip.getPqBits() / 8.0) + Float.BYTES) + + ip.getnLists() * Integer.BYTES); + return 2 * approximatedIvfBytes; + } else { + throw new IllegalArgumentException("Unsupported CAGRA build algo"); + } + } + + class ManagedCuVSResources { + + private final CuVSResources resource; + private final ReentrantLock lock; + private long neededMemory; + + public ManagedCuVSResources(CuVSResources resource) { + this.resource = resource; + lock = new ReentrantLock(); + } + + public CuVSResources getResource() { + return resource; + } + + public long getNeededMemory() { + return neededMemory; + } + + public void setNeededMemory(long neededMemory) { + this.neededMemory = neededMemory; + } + + public void lock() { + lock.lock(); + } + + public void unlock() { + lock.unlock(); + } + + public boolean isLocked() { + return lock.isLocked(); + } + } +} diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java deleted file mode 100644 index 07e64fb1..00000000 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSProvider.java +++ /dev/null @@ -1,168 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import com.nvidia.cuvs.BruteForceIndex; -import com.nvidia.cuvs.CagraIndex; -import com.nvidia.cuvs.CagraIndexParams; -import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; -import com.nvidia.cuvs.CagraIndexParams.HnswHeuristicType; -import com.nvidia.cuvs.CuVSDeviceMatrix; -import com.nvidia.cuvs.CuVSHostMatrix; -import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSMatrix.Builder; -import com.nvidia.cuvs.CuVSMatrix.DataType; -import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.GPUInfoProvider; -import com.nvidia.cuvs.HnswIndex; -import com.nvidia.cuvs.HnswIndexParams; -import com.nvidia.cuvs.TieredIndex; -import com.nvidia.cuvs.spi.CuVSProvider; -import java.lang.invoke.MethodHandle; -import java.nio.file.Path; -import java.util.logging.Level; - -class FilterCuVSProvider implements CuVSProvider { - - private final CuVSProvider delegate; - - FilterCuVSProvider(CuVSProvider delegate) { - this.delegate = delegate; - } - - @Override - public Path nativeLibraryPath() { - return CuVSProvider.TMPDIR; - } - - @Override - public CuVSResources newCuVSResources(Path tempPath) throws Throwable { - return delegate.newCuVSResources(tempPath); - } - - @Override - public BruteForceIndex.Builder newBruteForceIndexBuilder(CuVSResources cuVSResources) - throws UnsupportedOperationException { - return delegate.newBruteForceIndexBuilder(cuVSResources); - } - - @Override - public CagraIndex.Builder newCagraIndexBuilder(CuVSResources cuVSResources) - throws UnsupportedOperationException { - return delegate.newCagraIndexBuilder(cuVSResources); - } - - @Override - public HnswIndex.Builder newHnswIndexBuilder(CuVSResources cuVSResources) - throws UnsupportedOperationException { - return delegate.newHnswIndexBuilder(cuVSResources); - } - - @Override - public CagraIndex mergeCagraIndexes(CagraIndex[] arg0) throws Throwable { - return delegate.mergeCagraIndexes(arg0); - } - - @Override - public GPUInfoProvider gpuInfoProvider() { - return delegate.gpuInfoProvider(); - } - - @Override - public Builder newHostMatrixBuilder(long rows, long cols, DataType dataType) { - return delegate.newHostMatrixBuilder(rows, cols, dataType); - } - - @Override - public Builder newHostMatrixBuilder( - long rows, long cols, int maxRows, int maxCols, DataType dataType) { - return delegate.newHostMatrixBuilder(rows, cols, maxRows, maxCols, dataType); - } - - @Override - public Builder newDeviceMatrixBuilder( - CuVSResources resources, long rows, long cols, DataType dataType) { - return delegate.newDeviceMatrixBuilder(resources, rows, cols, dataType); - } - - @Override - public Builder newDeviceMatrixBuilder( - CuVSResources resources, long rows, long cols, int maxRows, int maxCols, DataType dataType) { - return delegate.newDeviceMatrixBuilder(resources, rows, cols, maxRows, maxCols, dataType); - } - - @Override - public MethodHandle newNativeMatrixBuilder() { - return delegate.newNativeMatrixBuilder(); - } - - @Override - public MethodHandle newNativeMatrixBuilderWithStrides() { - return delegate.newNativeMatrixBuilderWithStrides(); - } - - @Override - public CuVSMatrix newMatrixFromArray(float[][] vectors) { - return delegate.newMatrixFromArray(vectors); - } - - @Override - public CuVSMatrix newMatrixFromArray(int[][] vectors) { - return delegate.newMatrixFromArray(vectors); - } - - @Override - public CuVSMatrix newMatrixFromArray(byte[][] vectors) { - return delegate.newMatrixFromArray(vectors); - } - - @Override - public TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) - throws UnsupportedOperationException { - return delegate.newTieredIndexBuilder(cuVSResources); - } - - @Override - public CagraIndexParams cagraIndexParamsFromHnswParams( - long arg0, long arg1, int arg2, int arg3, HnswHeuristicType arg4, CuvsDistanceType arg5) { - return delegate.cagraIndexParamsFromHnswParams(arg0, arg1, arg2, arg3, arg4, arg5); - } - - @Override - public Level getLogLevel() { - return delegate.getLogLevel(); - } - - @Override - public void setLogLevel(Level arg0) { - delegate.setLogLevel(arg0); - } - - @Override - public HnswIndex hnswIndexFromCagra(HnswIndexParams arg0, CagraIndex arg1) throws Throwable { - return delegate.hnswIndexFromCagra(arg0, arg1); - } - - @Override - public void enableRMMManagedPooledMemory(int arg0, int arg1) { - delegate.enableRMMManagedPooledMemory(arg0, arg1); - } - - @Override - public void enableRMMPooledMemory(int arg0, int arg1) { - delegate.enableRMMPooledMemory(arg0, arg1); - } - - @Override - public void resetRMMPooledMemory() { - delegate.resetRMMPooledMemory(); - } - - @Override - public HnswIndex hnswIndexBuild(CuVSResources arg0, HnswIndexParams arg1, CuVSMatrix arg2) - throws Throwable { - return delegate.hnswIndexBuild(arg0, arg1, arg2); - } -} diff --git a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java b/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java deleted file mode 100644 index 83cebafa..00000000 --- a/src/main/java/com/nvidia/cuvs/lucene/FilterCuVSServiceProvider.java +++ /dev/null @@ -1,24 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import com.nvidia.cuvs.spi.CuVSProvider; -import com.nvidia.cuvs.spi.CuVSServiceProvider; - -/** - * A provider that creates instances of FilterCuVSProvider. - * - * @since 25.10 - */ -public class FilterCuVSServiceProvider extends CuVSServiceProvider { - - /** - * Initialize and return an CuVSProvider provided by this provider. - */ - @Override - public CuVSProvider get(CuVSProvider builtinProvider) { - return new FilterCuVSProvider(builtinProvider); - } -} diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index 3c42707c..df5e2941 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -31,6 +31,7 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; private static final int MAX_DIMENSIONS = 4096; private final AcceleratedHNSWParams acceleratedHNSWParams; + private final CuvsResourcesManager cuvsResourcesManager; static final String HNSW_META_CODEC_NAME = "Lucene99HnswVectorsFormatMeta"; static final String HNSW_META_CODEC_EXT = "vem"; @@ -67,6 +68,11 @@ public Lucene99AcceleratedHNSWVectorsFormat() { public Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams acceleratedHNSWParams) { super("Lucene99AcceleratedHNSWVectorsFormat"); this.acceleratedHNSWParams = acceleratedHNSWParams; + if (isSupported()) { + cuvsResourcesManager = new CuvsResourcesManager(acceleratedHNSWParams.getWriterThreads()); + } else { + cuvsResourcesManager = null; // Will not be needed in fallback mode. + } } /** @@ -77,7 +83,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException var flatWriter = FLAT_VECTORS_FORMAT.fieldsWriter(state); if (isSupported()) { log.log(Level.FINE, "cuVS is supported so using the Lucene99AcceleratedHNSWVectorsWriter"); - return new Lucene99AcceleratedHNSWVectorsWriter(state, acceleratedHNSWParams, flatWriter); + return new Lucene99AcceleratedHNSWVectorsWriter( + state, acceleratedHNSWParams, flatWriter, cuvsResourcesManager); } else { log.log( Level.WARNING, diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index bc79dfd3..a4d9cff6 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -16,7 +16,6 @@ import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_EXT; 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; @@ -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.CuvsResourcesManager.ManagedCuVSResources; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -62,6 +62,7 @@ public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; + private final CuvsResourcesManager cuvsResourcesManager; private IndexOutput hnswMeta = null; private IndexOutput hnswVectorIndex = null; private String vemFileName; @@ -88,12 +89,14 @@ public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { public Lucene99AcceleratedHNSWVectorsWriter( SegmentWriteState state, AcceleratedHNSWParams acceleratedHNSWParams, - FlatVectorsWriter flatVectorsWriter) + FlatVectorsWriter flatVectorsWriter, + CuvsResourcesManager cuvsResourcesManager) throws IOException { super(); this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; this.acceleratedHNSWParams = acceleratedHNSWParams; + this.cuvsResourcesManager = cuvsResourcesManager; vemFileName = IndexFileNames.segmentFileName( state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); @@ -145,8 +148,10 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException * @param fieldInfo instance of FieldInfo that has the field description * @param vectors vectors to index * @throws IOException + * @throws InterruptedException */ - private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) + throws IOException, InterruptedException { if (vectors.size() == 0) { writeEmpty(fieldInfo, hnswMeta); return; @@ -155,19 +160,21 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro writeSingleVectorGraph(fieldInfo, vectors); return; } + CagraIndexParams params = + cagraIndexParams( + acceleratedHNSWParams.getWriterThreads(), + acceleratedHNSWParams.getIntermediateGraphDegree(), + acceleratedHNSWParams.getGraphdegree(), + acceleratedHNSWParams.getCagraGraphBuildAlgo(), + acceleratedHNSWParams.getCuVSIvfPqParams()); + ManagedCuVSResources managedCuVSResources = + cuvsResourcesManager.acquireResource(vectors.size(), vectors.get(0).length, params); try { CuVSMatrix dataset = Utils.createFloatMatrix( - vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); - CagraIndexParams params = - cagraIndexParams( - acceleratedHNSWParams.getWriterThreads(), - acceleratedHNSWParams.getIntermediateGraphDegree(), - acceleratedHNSWParams.getGraphdegree(), - acceleratedHNSWParams.getCagraGraphBuildAlgo(), - acceleratedHNSWParams.getCuVSIvfPqParams()); + vectors, fieldInfo.getVectorDimension(), managedCuVSResources.getResource()); CagraIndex cagraIndex = - CagraIndex.newBuilder(getCuVSResourcesInstance()) + CagraIndex.newBuilder(managedCuVSResources.getResource()) .withDataset(dataset) .withIndexParams(params) .build(); @@ -184,7 +191,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro acceleratedHNSWParams.getHnswLayers(), acceleratedHNSWParams.getGraphdegree(), params, - QuantizationType.NONE); + QuantizationType.NONE, + managedCuVSResources.getResource()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; @@ -201,6 +209,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro cagraIndex.close(); } catch (Throwable t) { Utils.handleThrowable(t); + } finally { + cuvsResourcesManager.releaseResource(managedCuVSResources); } } @@ -211,10 +221,14 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro public void flush(int maxDoc, DocMap sortMap) throws IOException { flatVectorsWriter.flush(maxDoc, sortMap); for (var field : fields) { - if (sortMap == null) { - writeField(field); - } else { - writeSortingField(field, sortMap); + try { + if (sortMap == null) { + writeField(field); + } else { + writeSortingField(field, sortMap); + } + } catch (Exception e) { + throw new IOException(e.getMessage()); } } } @@ -224,8 +238,9 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { * * @param fieldData * @throws IOException + * @throws InterruptedException */ - private void writeField(FieldWriter fieldData) throws IOException { + private void writeField(FieldWriter fieldData) throws IOException, InterruptedException { writeFieldInternal(fieldData.fieldInfo(), fieldData.getFloatVectors()); } @@ -235,8 +250,10 @@ private void writeField(FieldWriter fieldData) throws IOException { * @param fieldData instance of GPUFieldWriter * @param sortMap instance of the DocMap * @throws IOException + * @throws InterruptedException */ - private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { + private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) + throws IOException, InterruptedException { DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; mapOldOrdToNewOrd(oldDocsWithFieldSet, sortMap, null, new2OldOrd, null); diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java index 0f8d9602..de2ce78a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java @@ -33,6 +33,7 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat extends KnnVector private static final int MAX_DIMENSIONS = 4096; private final AcceleratedHNSWParams acceleratedHNSWParams; + private final CuvsResourcesManager cuvsResourcesManager; static { try { @@ -63,6 +64,11 @@ public LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat( AcceleratedHNSWParams acceleratedHNSWParams) { super("Lucene99AcceleratedHNSWBinaryQuantizedVectorsFormat"); this.acceleratedHNSWParams = acceleratedHNSWParams; + if (isSupported()) { + cuvsResourcesManager = new CuvsResourcesManager(acceleratedHNSWParams.getWriterThreads()); + } else { + cuvsResourcesManager = null; // Will not be needed in fallback mode. + } } /** @@ -76,7 +82,7 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException Level.FINE, "cuVS is supported so using the Lucene99AcceleratedHNSWBinaryQuantizedVectorsWriter"); return new LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter( - state, acceleratedHNSWParams, flatWriter); + state, acceleratedHNSWParams, flatWriter, cuvsResourcesManager); } else { try { // Fallback to Lucene's Lucene102HnswBinaryQuantizedVectorsFormat format diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index ade0e702..0a0d8a02 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -16,8 +16,6 @@ import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_EXT; import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_EXT; 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 org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; @@ -26,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.CuvsResourcesManager.ManagedCuVSResources; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -65,6 +64,7 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter extends KnnVector private final List fields = new ArrayList<>(); private final InfoStream infoStream; private final AcceleratedHNSWParams acceleratedHNSWParams; + private final CuvsResourcesManager cuvsResourcesManager; private IndexOutput hnswMeta = null, hnswVectorIndex = null; private boolean finished; private String vemFileName; @@ -81,13 +81,14 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter extends KnnVector public LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter( SegmentWriteState state, AcceleratedHNSWParams acceleratedHNSWParams, - FlatVectorsWriter flatVectorsWriter) + FlatVectorsWriter flatVectorsWriter, + CuvsResourcesManager cuvsResourcesManager) throws IOException { super(); this.acceleratedHNSWParams = acceleratedHNSWParams; this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; - + this.cuvsResourcesManager = cuvsResourcesManager; vemFileName = IndexFileNames.segmentFileName( state.segmentInfo.name, state.segmentSuffix, HNSW_META_CODEC_EXT); @@ -145,34 +146,39 @@ public KnnFieldVectorsWriter addField(FieldInfo fieldInfo) throws IOException * @param fieldInfo instance of FieldInfo that has the field description * @param vectors binary quantized vectors (packed bits as bytes) * @throws IOException + * @throws InterruptedException */ - private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) + throws IOException, InterruptedException { if (vectors.size() == 0) { writeEmpty(fieldInfo, hnswMeta); return; } + CagraIndexParams params = + cagraIndexParams( + acceleratedHNSWParams.getWriterThreads(), + acceleratedHNSWParams.getIntermediateGraphDegree(), + acceleratedHNSWParams.getGraphdegree(), + acceleratedHNSWParams.getCagraGraphBuildAlgo(), + acceleratedHNSWParams.getCuVSIvfPqParams()); + ManagedCuVSResources managedCuVSResources = + cuvsResourcesManager.acquireResource(vectors.size(), vectors.get(0).length, params); + try { int dimensions = fieldInfo.getVectorDimension(); int bytesPerVector = (dimensions + 7) / 8; CuVSMatrix dataset = - Utils.createByteMatrix(vectors, bytesPerVector, getCuVSResourcesInstance()); + Utils.createByteMatrix(vectors, bytesPerVector, managedCuVSResources.getResource()); if (dataset.size() < 2) { writeSingleVectorGraph(fieldInfo, vectors); return; } - CagraIndexParams params = - cagraIndexParams( - acceleratedHNSWParams.getWriterThreads(), - acceleratedHNSWParams.getIntermediateGraphDegree(), - acceleratedHNSWParams.getGraphdegree(), - acceleratedHNSWParams.getCagraGraphBuildAlgo(), - acceleratedHNSWParams.getCuVSIvfPqParams()); CagraIndex cagraIndex = - CagraIndex.newBuilder(getCuVSResourcesInstance()) + CagraIndex.newBuilder(managedCuVSResources.getResource()) .withDataset(dataset) .withIndexParams(params) .build(); @@ -191,7 +197,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw acceleratedHNSWParams.getHnswLayers(), acceleratedHNSWParams.getGraphdegree(), params, - QuantizationType.BINARY); + QuantizationType.BINARY, + managedCuVSResources.getResource()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index @@ -214,6 +221,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw } catch (Throwable t) { Utils.handleThrowable(t); + } finally { + cuvsResourcesManager.releaseResource(managedCuVSResources); } } @@ -223,12 +232,16 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw @Override public void flush(int maxDoc, DocMap sortMap) throws IOException { flatVectorsWriter.flush(maxDoc, sortMap); - for (var field : fields) { - if (sortMap == null) { - writeField(field); - } else { - writeSortingField(field, sortMap); + try { + for (var field : fields) { + if (sortMap == null) { + writeField(field); + } else { + writeSortingField(field, sortMap); + } } + } catch (Exception e) { + throw new IOException(e.getMessage()); } } @@ -237,8 +250,9 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { * * @param fieldData * @throws IOException + * @throws InterruptedException */ - private void writeField(FieldWriter fieldData) throws IOException { + private void writeField(FieldWriter fieldData) throws IOException, InterruptedException { writeFieldInternal(fieldData.fieldInfo(), fieldData.getByteVectors()); } @@ -248,8 +262,10 @@ private void writeField(FieldWriter fieldData) throws IOException { * @param fieldData instance of BinaryQuantizedGPUFieldWriter * @param sortMap instance of the DocMap * @throws IOException + * @throws InterruptedException */ - private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { + private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) + throws IOException, InterruptedException { DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord @@ -362,7 +378,6 @@ public void finish() throws IOException { @Override public void close() throws IOException { IOUtils.close(hnswMeta, hnswVectorIndex, flatVectorsWriter); - closeCuVSResourcesInstance(); } /** diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java index 534390b8..7cfb083a 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java @@ -30,6 +30,7 @@ public class LuceneAcceleratedHNSWScalarQuantizedVectorsFormat extends KnnVector private static final int MAX_DIMENSIONS = 4096; private final AcceleratedHNSWParams acceleratedHNSWParams; + private final CuvsResourcesManager cuvsResourcesManager; static { try { @@ -58,6 +59,11 @@ public LuceneAcceleratedHNSWScalarQuantizedVectorsFormat( AcceleratedHNSWParams acceleratedHNSWParams) { super("Lucene99AcceleratedHNSWScalarQuantizedVectorsFormat"); this.acceleratedHNSWParams = acceleratedHNSWParams; + if (isSupported()) { + cuvsResourcesManager = new CuvsResourcesManager(acceleratedHNSWParams.getWriterThreads()); + } else { + cuvsResourcesManager = null; // Will not be needed in fallback mode. + } } /** @@ -69,7 +75,7 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException if (isSupported()) { log.info("cuVS is supported so using the Lucene99AcceleratedHNSWQuantizedVectorsWriter"); return new LuceneAcceleratedHNSWScalarQuantizedVectorsWriter( - state, acceleratedHNSWParams, flatWriter); + state, acceleratedHNSWParams, flatWriter, cuvsResourcesManager); } else { try { // Fallback to Lucene's Lucene99HnswScalarQuantizedVectorsFormat diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 3e3713df..11a11b83 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -16,8 +16,6 @@ import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_EXT; import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_EXT; 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 org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; @@ -26,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.CuvsResourcesManager.ManagedCuVSResources; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -66,6 +65,7 @@ public class LuceneAcceleratedHNSWScalarQuantizedVectorsWriter extends KnnVector private final List fields = new ArrayList<>(); private final InfoStream infoStream; private final AcceleratedHNSWParams acceleratedHNSWParams; + private final CuvsResourcesManager cuvsResourcesManager; private IndexOutput hnswMeta = null, hnswVectorIndex = null; private boolean finished; private String vemFileName; @@ -91,12 +91,14 @@ public class LuceneAcceleratedHNSWScalarQuantizedVectorsWriter extends KnnVector public LuceneAcceleratedHNSWScalarQuantizedVectorsWriter( SegmentWriteState state, AcceleratedHNSWParams acceleratedHNSWParams, - FlatVectorsWriter flatVectorsWriter) + FlatVectorsWriter flatVectorsWriter, + CuvsResourcesManager cuvsResourcesManager) throws IOException { super(); this.acceleratedHNSWParams = acceleratedHNSWParams; this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; + this.cuvsResourcesManager = cuvsResourcesManager; vemFileName = IndexFileNames.segmentFileName( @@ -165,13 +167,26 @@ private static byte[] convertSignedToUnsigned(byte[] signedVector) { * @param fieldInfo instance of FieldInfo that has the field description * @param vectors quantized vectors * @throws IOException + * @throws InterruptedException */ - private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOException { + private void writeFieldInternal(FieldInfo fieldInfo, List vectors) + throws IOException, InterruptedException { if (vectors.size() == 0) { writeEmpty(fieldInfo, hnswMeta); return; } + CagraIndexParams params = + cagraIndexParams( + acceleratedHNSWParams.getWriterThreads(), + acceleratedHNSWParams.getIntermediateGraphDegree(), + acceleratedHNSWParams.getGraphdegree(), + acceleratedHNSWParams.getCagraGraphBuildAlgo(), + acceleratedHNSWParams.getCuVSIvfPqParams()); + ManagedCuVSResources managedCuVSResources = + cuvsResourcesManager.acquireResource( + vectors.size(), fieldInfo.getVectorDimension(), params); + try { int dimensions = fieldInfo.getVectorDimension(); @@ -183,22 +198,15 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE // Create CuVSMatrix with BYTE data type (unsigned bytes) CuVSMatrix dataset = - Utils.createByteMatrix(unsignedVectors, dimensions, getCuVSResourcesInstance()); + Utils.createByteMatrix(unsignedVectors, dimensions, managedCuVSResources.getResource()); if (dataset.size() < 2) { writeSingleVectorGraph(fieldInfo, unsignedVectors); return; } - CagraIndexParams params = - cagraIndexParams( - acceleratedHNSWParams.getWriterThreads(), - acceleratedHNSWParams.getIntermediateGraphDegree(), - acceleratedHNSWParams.getGraphdegree(), - acceleratedHNSWParams.getCagraGraphBuildAlgo(), - acceleratedHNSWParams.getCuVSIvfPqParams()); CagraIndex cagraIndex = - CagraIndex.newBuilder(getCuVSResourcesInstance()) + CagraIndex.newBuilder(managedCuVSResources.getResource()) .withDataset(dataset) .withIndexParams(params) .build(); @@ -216,7 +224,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE acceleratedHNSWParams.getHnswLayers(), acceleratedHNSWParams.getGraphdegree(), params, - QuantizationType.SCALAR); + QuantizationType.SCALAR, + managedCuVSResources.getResource()); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); @@ -240,6 +249,8 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE cagraIndex.close(); } catch (Throwable t) { Utils.handleThrowable(t); + } finally { + cuvsResourcesManager.releaseResource(managedCuVSResources); } } @@ -249,12 +260,16 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE @Override public void flush(int maxDoc, DocMap sortMap) throws IOException { flatVectorsWriter.flush(maxDoc, sortMap); - for (var field : fields) { - if (sortMap == null) { - writeField(field); - } else { - writeSortingField(field, sortMap); + try { + for (var field : fields) { + if (sortMap == null) { + writeField(field); + } else { + writeSortingField(field, sortMap); + } } + } catch (Exception e) { + throw new IOException(e.getMessage()); } } @@ -263,8 +278,9 @@ public void flush(int maxDoc, DocMap sortMap) throws IOException { * * @param fieldData * @throws IOException + * @throws InterruptedException */ - private void writeField(FieldWriter fieldData) throws IOException { + private void writeField(FieldWriter fieldData) throws IOException, InterruptedException { writeFieldInternal(fieldData.fieldInfo(), fieldData.getByteVectors()); } @@ -274,8 +290,10 @@ private void writeField(FieldWriter fieldData) throws IOException { * @param fieldData instance of ScalarQuantizedGPUFieldWriter * @param sortMap instance of the DocMap * @throws IOException + * @throws InterruptedException */ - private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) throws IOException { + private void writeSortingField(FieldWriter fieldData, Sorter.DocMap sortMap) + throws IOException, InterruptedException { DocsWithFieldSet oldDocsWithFieldSet = fieldData.getDocsWithFieldSet(); final int[] new2OldOrd = new int[oldDocsWithFieldSet.cardinality()]; // new ord to old ord @@ -386,7 +404,6 @@ public void finish() throws IOException { @Override public void close() throws IOException { IOUtils.close(hnswMeta, hnswVectorIndex, flatVectorsWriter); - closeCuVSResourcesInstance(); } /** From f7ba11616097dcce31262ec9f91f860625a3ef97 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Tue, 28 Apr 2026 22:19:28 -0400 Subject: [PATCH 2/8] Variable renaming and simplification in CuvsResourcesManager --- .../cuvs/lucene/CuvsResourcesManager.java | 78 +++++++------------ 1 file changed, 28 insertions(+), 50 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java b/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java index fc80e8cd..372145a7 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java @@ -14,6 +14,8 @@ import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.GPUInfoProvider; import com.nvidia.cuvs.spi.CuVSProvider; +import java.util.Arrays; +import java.util.Objects; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.ReentrantLock; @@ -22,7 +24,7 @@ public class CuvsResourcesManager { - private static final Logger log = Logger.getLogger(Utils.class.getName()); + private static final Logger LOG = Logger.getLogger(Utils.class.getName()); private static final CuVSProvider PROVIDER = CuVSProvider.provider(); private static final GPUInfoProvider GPU_INFO_PROVIDER = PROVIDER.gpuInfoProvider(); private static final int MAX_POOL_SIZE = 512; @@ -53,56 +55,35 @@ public CuvsResourcesManager(int capacity) { } private ManagedCuVSResources getAvailableResourcesFromPool() { - try { - lock.lock(); - for (int i = 0; i < capacity; i++) { - if (pool[i] != null && !pool[i].isLocked()) { - return pool[i]; - } - } - } finally { - lock.unlock(); - } - return null; + return Arrays.stream(pool).filter(mcr -> !mcr.isLocked()).findFirst().orElse(null); } - private int getNumLockedResources() { - try { - lock.lock(); - int res = 0; - for (int i = 0; i < capacity; i++) { - if (pool[i].isLocked()) { - res += 1; - } - } - return res; - } finally { - lock.unlock(); - } + private long getNumLockedResources() { + return Arrays.stream(pool).filter(mcr -> mcr.isLocked()).count(); } public ManagedCuVSResources acquireResource(long rows, long dimension, CagraIndexParams params) throws InterruptedException { try { lock.lock(); - long neededMem = getEstimatedMemoryRequirement(rows, dimension, params); + long neededMemory = getEstimatedMemoryRequirement(rows, dimension, params); - if (neededMem > totalDeviceMemory) { + if (neededMemory > totalDeviceMemory) { throw new RuntimeException("Not enough GPU device memory available"); } while (getNumLockedResources() == capacity - || (totalDeviceMemory - reserveMemory.get()) < neededMem) { + || (totalDeviceMemory - reserveMemory.get()) < neededMemory) { resourcesAvailable.await(); } - reserveMemory.addAndGet(neededMem); + reserveMemory.addAndGet(neededMemory); - ManagedCuVSResources res = getAvailableResourcesFromPool(); - assert res != null; + ManagedCuVSResources managedCuVSResources = getAvailableResourcesFromPool(); + assert managedCuVSResources != null; - res.setNeededMemory(neededMem); - res.lock(); - return res; + managedCuVSResources.setNeededMemory(neededMemory); + managedCuVSResources.lock(); + return managedCuVSResources; } finally { lock.unlock(); } @@ -120,30 +101,27 @@ public void releaseResource(ManagedCuVSResources resource) { } public void shutdown() { - try { - lock.lock(); - for (int i = 0; i < capacity; i++) { - if (pool[i] != null && !pool[i].isLocked() && pool[i].getResource() != null) { - pool[i].getResource().close(); - } - } - } finally { - lock.unlock(); - } + Arrays.stream(pool) + .forEach( + mcr -> { + if (Objects.nonNull(mcr) && Objects.nonNull(mcr.getResource())) { + mcr.getResource().close(); + } + }); } private static CuVSResources getCuVSResourceInstance() { try { return CuVSResources.create(); } catch (UnsupportedOperationException uoe) { - log.log( + LOG.log( Level.WARNING, "cuVS is not supported on this platform or java version: " + uoe.getMessage()); } catch (Throwable t) { if (t instanceof ExceptionInInitializerError ex) { t = ex.getCause(); } - log.log(Level.WARNING, "Exception occurred during creation of cuVS resources. " + t); + LOG.log(Level.WARNING, "Exception occurred during creation of cuVS resources. " + t); } return null; } @@ -166,17 +144,17 @@ private long getEstimatedMemoryRequirement(long rows, long dimension, CagraIndex class ManagedCuVSResources { - private final CuVSResources resource; + private final CuVSResources cuVSResources; private final ReentrantLock lock; private long neededMemory; - public ManagedCuVSResources(CuVSResources resource) { - this.resource = resource; + public ManagedCuVSResources(CuVSResources cuVSResources) { + this.cuVSResources = cuVSResources; lock = new ReentrantLock(); } public CuVSResources getResource() { - return resource; + return cuVSResources; } public long getNeededMemory() { From 0cf52db176485f4add1e15550103848eff20dbca Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Tue, 28 Apr 2026 23:43:42 -0400 Subject: [PATCH 3/8] Cleanup --- .../com/nvidia/cuvs/lucene/CuvsResourcesManager.java | 10 +++++++--- .../lucene/Lucene99AcceleratedHNSWVectorsWriter.java | 2 -- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java b/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java index 372145a7..2c9c5dc9 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java @@ -131,17 +131,21 @@ private long getEstimatedMemoryRequirement(long rows, long dimension, CagraIndex if (buildAlgo.equals(NN_DESCENT)) { return 2 * rows * dimension * Float.BYTES; } else if (buildAlgo.equals(IVF_PQ)) { + assert params.getCuVSIvfPqParams() != null; CuVSIvfPqIndexParams ip = params.getCuVSIvfPqParams().getIndexParams(); - long approximatedIvfBytes = - (long) + assert ip != null; + return 2 + * (long) (rows * (ip.getPqDim() * (ip.getPqBits() / 8.0) + Float.BYTES) + ip.getnLists() * Integer.BYTES); - return 2 * approximatedIvfBytes; } else { throw new IllegalArgumentException("Unsupported CAGRA build algo"); } } + /** + * Holds reference to CuVSResources with its associated lock, and needed memory. + */ class ManagedCuVSResources { private final CuVSResources cuVSResources; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index a4d9cff6..680843da 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -15,7 +15,6 @@ import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_EXT; import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_EXT; 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.Utils.createListFromMergedVectors; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; @@ -346,7 +345,6 @@ public void finish() throws IOException { public void close() throws IOException { printInfoStream(infoStream, COMPONENT, "Closing resources"); IOUtils.close(hnswMeta, hnswVectorIndex, flatVectorsWriter); - closeCuVSResourcesInstance(); } /** From 5d331d9989e8566aa84dd1ec66c5307b7939f4b9 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Wed, 29 Apr 2026 21:33:59 -0400 Subject: [PATCH 4/8] Add javadocs and rename variables --- .../cuvs/lucene/CuvsResourcesManager.java | 63 ++++++++++++++----- 1 file changed, 48 insertions(+), 15 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java b/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java index 2c9c5dc9..276b058b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java @@ -22,6 +22,10 @@ import java.util.logging.Level; import java.util.logging.Logger; +/** + * Manages a pool of finite {@link ManagedCuVSResources} and allows for the accessing threads + * to lock and acquire available instance and release them back to the pool when finished. + */ public class CuvsResourcesManager { private static final Logger LOG = Logger.getLogger(Utils.class.getName()); @@ -48,20 +52,22 @@ public CuvsResourcesManager(int capacity) { for (int i = 0; i < capacity; i++) { pool[i] = new ManagedCuVSResources(getCuVSResourceInstance()); } - CuVSResources rx = getCuVSResourceInstance(); + CuVSResources cuVSResources = getCuVSResourceInstance(); reserveMemory = new AtomicLong(); - totalDeviceMemory = GPU_INFO_PROVIDER.getCurrentInfo(rx).totalDeviceMemoryInBytes(); - rx.close(); - } - - private ManagedCuVSResources getAvailableResourcesFromPool() { - return Arrays.stream(pool).filter(mcr -> !mcr.isLocked()).findFirst().orElse(null); - } - - private long getNumLockedResources() { - return Arrays.stream(pool).filter(mcr -> mcr.isLocked()).count(); + totalDeviceMemory = GPU_INFO_PROVIDER.getCurrentInfo(cuVSResources).totalDeviceMemoryInBytes(); + cuVSResources.close(); } + /** + * Acquire an instance of {@link ManagedCuVSResources} when available and enough + * device memory is also available for the request to complete. + * + * @param rows the number of vectors in the dataset + * @param dimension the vector dimension in the dataset + * @param params an instance of {@link CagraIndexParams} + * @return an instance of {@link ManagedCuVSResources} + * @throws InterruptedException + */ public ManagedCuVSResources acquireResource(long rows, long dimension, CagraIndexParams params) throws InterruptedException { try { @@ -72,7 +78,7 @@ public ManagedCuVSResources acquireResource(long rows, long dimension, CagraInde throw new RuntimeException("Not enough GPU device memory available"); } - while (getNumLockedResources() == capacity + while (getNumberOfUnavailableResources() == capacity || (totalDeviceMemory - reserveMemory.get()) < neededMemory) { resourcesAvailable.await(); } @@ -89,10 +95,16 @@ public ManagedCuVSResources acquireResource(long rows, long dimension, CagraInde } } + /** + * Releases the acquired instance of {@link ManagedCuVSResources} back to the pool. + * + * @param resource the acquired instance of {@link ManagedCuVSResources} by the thread + */ public void releaseResource(ManagedCuVSResources resource) { try { lock.lock(); reserveMemory.addAndGet(-resource.getNeededMemory()); + resource.resetNeededMemory(); resource.unlock(); resourcesAvailable.signalAll(); } finally { @@ -100,12 +112,16 @@ public void releaseResource(ManagedCuVSResources resource) { } } + /** + * Shuts down the instances of wrapped {@link CuVSResources} in the pool. + */ public void shutdown() { Arrays.stream(pool) .forEach( - mcr -> { - if (Objects.nonNull(mcr) && Objects.nonNull(mcr.getResource())) { - mcr.getResource().close(); + managedCuVSResources -> { + if (Objects.nonNull(managedCuVSResources) + && Objects.nonNull(managedCuVSResources.getResource())) { + managedCuVSResources.getResource().close(); } }); } @@ -126,6 +142,19 @@ private static CuVSResources getCuVSResourceInstance() { return null; } + private ManagedCuVSResources getAvailableResourcesFromPool() { + return Arrays.stream(pool) + .filter(managedCuVSResources -> !managedCuVSResources.isLocked()) + .findFirst() + .orElse(null); + } + + private long getNumberOfUnavailableResources() { + return Arrays.stream(pool) + .filter(managedCuVSResources -> managedCuVSResources.isLocked()) + .count(); + } + private long getEstimatedMemoryRequirement(long rows, long dimension, CagraIndexParams params) { CagraGraphBuildAlgo buildAlgo = params.getCagraGraphBuildAlgo(); if (buildAlgo.equals(NN_DESCENT)) { @@ -165,6 +194,10 @@ public long getNeededMemory() { return neededMemory; } + public void resetNeededMemory() { + setNeededMemory(0); + } + public void setNeededMemory(long neededMemory) { this.neededMemory = neededMemory; } From 89d399189010d8ad682df62af3785f6bc4f6bca2 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Mon, 11 May 2026 17:33:15 -0400 Subject: [PATCH 5/8] Add peak device memory estimation methods for NN_DESCENT and IVF_PQ --- .../cuvs/lucene/CuvsResourcesManager.java | 84 +++++++++++++++++-- 1 file changed, 77 insertions(+), 7 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java b/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java index 276b058b..a5fc58e2 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java @@ -11,6 +11,8 @@ import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; import com.nvidia.cuvs.CuVSIvfPqIndexParams; +import com.nvidia.cuvs.CuVSIvfPqParams; +import com.nvidia.cuvs.CuVSIvfPqSearchParams; import com.nvidia.cuvs.CuVSResources; import com.nvidia.cuvs.GPUInfoProvider; import com.nvidia.cuvs.spi.CuVSProvider; @@ -21,6 +23,7 @@ import java.util.concurrent.locks.ReentrantLock; import java.util.logging.Level; import java.util.logging.Logger; +import java.util.stream.LongStream; /** * Manages a pool of finite {@link ManagedCuVSResources} and allows for the accessing threads @@ -158,20 +161,87 @@ private long getNumberOfUnavailableResources() { private long getEstimatedMemoryRequirement(long rows, long dimension, CagraIndexParams params) { CagraGraphBuildAlgo buildAlgo = params.getCagraGraphBuildAlgo(); if (buildAlgo.equals(NN_DESCENT)) { - return 2 * rows * dimension * Float.BYTES; + return estimateNNDescentIndexBuildPeakMemory(rows, dimension, params); } else if (buildAlgo.equals(IVF_PQ)) { assert params.getCuVSIvfPqParams() != null; - CuVSIvfPqIndexParams ip = params.getCuVSIvfPqParams().getIndexParams(); - assert ip != null; - return 2 - * (long) - (rows * (ip.getPqDim() * (ip.getPqBits() / 8.0) + Float.BYTES) - + ip.getnLists() * Integer.BYTES); + return estimateIVFPQIndexBuildPeakMemory(rows, dimension, params); } else { throw new IllegalArgumentException("Unsupported CAGRA build algo"); } } + private long estimateNNDescentIndexBuildPeakMemory( + long rows, long dimension, CagraIndexParams params) { + /* + * N = Number of vectors + * D = Vector dimension + * I = Intermediate Graph Degree + * Sidx = Bytes per graph neighbor ID. This is sizeof(IdxT), usually 4 for int32_t or uint32_t. + * + * NND_device_peak = N × (D × 2 + 276) + * optimize_peak = N × (4 + (Sidx + 1) × I) + * build_peak = dataset_size + max(NND_device_peak, optimize_peak) + */ + + final int sIdx = 4; + long datasetSize = rows * dimension * Float.BYTES; + long nnDevicePeak = rows * (dimension * 2 + 276); + long optimizePeak = rows * (4 + (sIdx + 1) * params.getIntermediateGraphDegree()); + return datasetSize + Math.max(nnDevicePeak, optimizePeak); + } + + private long estimateIVFPQIndexBuildPeakMemory( + long rows, long dimension, CagraIndexParams params) { + + /* + * R = IVF-PQ training-set ratio. This is train_set_ratio. + * N = Number of vectors + * D = Vector dimension + * C = Number of IVF-PQ coarse clusters/lists. This is the IVF-PQ n_lists + * Q = Query batch size + * I = Intermediate graph degree + * Sidx = Bytes per graph neighbor ID. This is sizeof(IdxT), usually 4 for int32_t or uint32_t. + * + * IVFPQ_build_peak ​= (R/N × D × 4) + (C × D × 4) + (R/N * sizeof(uint32_t)) + * IVFPQ_search_peak = (Q × D × 4)+ (Q × I × sizeof(uint32_t))+ (Q × I × 4) + * optimize_peak = N × (4 + (Sidx + 1) × I) + * build_peak = dataset_size + max(IVFPQ_build_peak, IVFPQ_search_peak, optimize_peak) + * + * https://github.com/rapidsai/cuvs/blob/main/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh#L1253 + * trainset_ratio = max(1, n_rows / max(kmeans_trainset_fraction * n_rows, n_lists)) + */ + + CuVSIvfPqParams p = params.getCuVSIvfPqParams(); + CuVSIvfPqIndexParams ip = p.getIndexParams(); + CuVSIvfPqSearchParams sp = p.getSearchParams(); + assert ip != null; + assert sp != null; + + final int sIdx = 4; + long datasetSize = rows * dimension * Float.BYTES; + + double trainsetRatio = + Math.max(1, rows / Math.max((ip.getKmeansTrainsetFraction() * rows), ip.getnLists())); + + double ivfPQBuildPeak = + (trainsetRatio / rows * dimension * 4) + + (ip.getnLists() * dimension * 4) + + (trainsetRatio / rows * 4); + + long queryBatchSize = 4096; // Not sure about this yet "max_internal_batch_size" + + long ivfPQSearchPeak = + (queryBatchSize * dimension * 4) + + (2 * queryBatchSize * params.getIntermediateGraphDegree() * 4); + + long optimizePeak = rows * (4 + (sIdx + 1) * params.getIntermediateGraphDegree()); + + return datasetSize + + LongStream.of((long) Math.ceil(ivfPQBuildPeak), ivfPQSearchPeak, optimizePeak) + .max() + .getAsLong(); + } + /** * Holds reference to CuVSResources with its associated lock, and needed memory. */ From a146fdc9e76cf7120023ed36c00c7e127d74c189 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Tue, 12 May 2026 00:05:31 -0400 Subject: [PATCH 6/8] Add logic to check actual current free device memory, change algo threshold value to 1M, and tweak estimations --- .../cuvs/lucene/CagraIndexParamsFactory.java | 2 +- .../cuvs/lucene/CuvsResourcesManager.java | 25 ++++++++++++------- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java b/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java index 20fed04b..8995460b 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CagraIndexParamsFactory.java @@ -19,7 +19,7 @@ */ public class CagraIndexParamsFactory { - private static final int ALGO_SWITCH_THRESHOLD = 5_000_000; + private static final int ALGO_SWITCH_THRESHOLD = 1_000_000; /** * Translation of the internal logic found here: diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java b/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java index a5fc58e2..fc8d822c 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java @@ -42,6 +42,7 @@ public class CuvsResourcesManager { private AtomicLong reserveMemory; private long totalDeviceMemory; private int capacity; + private CuVSResources cuVSResources; public CuvsResourcesManager(int capacity) { if (capacity > MAX_POOL_SIZE || capacity <= 0) { @@ -55,10 +56,9 @@ public CuvsResourcesManager(int capacity) { for (int i = 0; i < capacity; i++) { pool[i] = new ManagedCuVSResources(getCuVSResourceInstance()); } - CuVSResources cuVSResources = getCuVSResourceInstance(); + cuVSResources = getCuVSResourceInstance(); reserveMemory = new AtomicLong(); totalDeviceMemory = GPU_INFO_PROVIDER.getCurrentInfo(cuVSResources).totalDeviceMemoryInBytes(); - cuVSResources.close(); } /** @@ -81,8 +81,12 @@ public ManagedCuVSResources acquireResource(long rows, long dimension, CagraInde throw new RuntimeException("Not enough GPU device memory available"); } + long currentFreeMemory = + GPU_INFO_PROVIDER.getCurrentInfo(cuVSResources).freeDeviceMemoryInBytes(); + while (getNumberOfUnavailableResources() == capacity - || (totalDeviceMemory - reserveMemory.get()) < neededMemory) { + || (totalDeviceMemory - reserveMemory.get()) < neededMemory + || currentFreeMemory < neededMemory) { resourcesAvailable.await(); } reserveMemory.addAndGet(neededMemory); @@ -127,6 +131,9 @@ public void shutdown() { managedCuVSResources.getResource().close(); } }); + if (cuVSResources != null) { + cuVSResources.close(); + } } private static CuVSResources getCuVSResourceInstance() { @@ -187,7 +194,7 @@ private long estimateNNDescentIndexBuildPeakMemory( long datasetSize = rows * dimension * Float.BYTES; long nnDevicePeak = rows * (dimension * 2 + 276); long optimizePeak = rows * (4 + (sIdx + 1) * params.getIntermediateGraphDegree()); - return datasetSize + Math.max(nnDevicePeak, optimizePeak); + return (long) ((datasetSize + Math.max(nnDevicePeak, optimizePeak)) * 1.25); } private long estimateIVFPQIndexBuildPeakMemory( @@ -218,7 +225,6 @@ private long estimateIVFPQIndexBuildPeakMemory( assert sp != null; final int sIdx = 4; - long datasetSize = rows * dimension * Float.BYTES; double trainsetRatio = Math.max(1, rows / Math.max((ip.getKmeansTrainsetFraction() * rows), ip.getnLists())); @@ -236,10 +242,11 @@ private long estimateIVFPQIndexBuildPeakMemory( long optimizePeak = rows * (4 + (sIdx + 1) * params.getIntermediateGraphDegree()); - return datasetSize - + LongStream.of((long) Math.ceil(ivfPQBuildPeak), ivfPQSearchPeak, optimizePeak) - .max() - .getAsLong(); + return (long) + (LongStream.of((long) Math.ceil(ivfPQBuildPeak), ivfPQSearchPeak, optimizePeak) + .max() + .getAsLong() + * 1.25); } /** From 64e0e520af42895db9b740baee4e420700177e76 Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Wed, 20 May 2026 22:04:19 -0400 Subject: [PATCH 7/8] Update utils - allow to switch between building matrix on device and host --- .../cuvs/lucene/CuVS2510GPUVectorsWriter.java | 5 ++-- .../Lucene99AcceleratedHNSWVectorsWriter.java | 3 +- .../java/com/nvidia/cuvs/lucene/Utils.java | 28 ++++++++++++------- 3 files changed, 23 insertions(+), 13 deletions(-) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java index e8fb304e..a4498d95 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsWriter.java @@ -11,6 +11,7 @@ import static com.nvidia.cuvs.lucene.CuVS2510GPUVectorsFormat.VERSION_CURRENT; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.closeCuVSResourcesInstance; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.getCuVSResourcesInstance; +import static com.nvidia.cuvs.lucene.Utils.Target.HOST; import static com.nvidia.cuvs.lucene.Utils.info; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; @@ -202,7 +203,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro try { CuVSMatrix cagraDataset = Utils.createFloatMatrix( - vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance(), HOST); writeCagraIndex(cagraIndexOutputStream, cagraDataset); } catch (Throwable t) { // Fallback to brute force in a few cases, for now. @@ -216,7 +217,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro var bruteForceIndexOutputStream = new IndexOutputOutputStream(cuvsIndex); CuVSMatrix bruteforceDataset = Utils.createFloatMatrix( - vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance()); + vectors, fieldInfo.getVectorDimension(), getCuVSResourcesInstance(), HOST); writeBruteForceIndex(bruteForceIndexOutputStream, bruteforceDataset); bruteForceIndexLength = cuvsIndex.getFilePointer() - bruteForceIndexOffset; diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index cee3823c..9465c311 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -14,6 +14,7 @@ import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_INDEX_EXT; import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_EXT; import static com.nvidia.cuvs.lucene.Lucene99AcceleratedHNSWVectorsFormat.HNSW_META_CODEC_NAME; +import static com.nvidia.cuvs.lucene.Utils.Target.HOST; import static com.nvidia.cuvs.lucene.Utils.createListFromMergedVectors; import static org.apache.lucene.index.VectorEncoding.FLOAT32; import static org.apache.lucene.util.RamUsageEstimator.shallowSizeOfInstance; @@ -166,7 +167,7 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) try { CuVSMatrix dataset = Utils.createFloatMatrix( - vectors, fieldInfo.getVectorDimension(), managedCuVSResources.getResource()); + vectors, fieldInfo.getVectorDimension(), managedCuVSResources.getResource(), HOST); CagraIndex cagraIndex = CagraIndex.newBuilder(managedCuVSResources.getResource()) .withDataset(dataset) diff --git a/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/src/main/java/com/nvidia/cuvs/lucene/Utils.java index c55aa7c3..28853e44 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -27,6 +27,11 @@ public class Utils { static final Logger log = Logger.getLogger(Utils.class.getName()); + public enum Target { + DEVICE, + HOST + } + /** * A utility method that throws specific types of throwable objects based on types. * @@ -51,19 +56,22 @@ static void handleThrowable(Throwable t) throws IOException { * @param data The float vectors * @param dimensions The number float elements in each vector * @param resources The CuVS resources for device matrix creation + * @param target To build the matrix on device or host * @return an instance of 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 - CuVSMatrix.Builder builder = - CuVSMatrix.deviceBuilder( - resources, - data.size(), // rows (number of vectors) - dimensions, // columns (vector dimension) - CuVSMatrix.DataType.FLOAT); + static CuVSMatrix createFloatMatrix( + List data, int dimensions, CuVSResources resources, Target target) { + CuVSMatrix.Builder builder = null; + + switch (target) { + case DEVICE: + builder = + CuVSMatrix.deviceBuilder(resources, data.size(), dimensions, CuVSMatrix.DataType.FLOAT); + case HOST: + builder = CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT); + } + assert builder != null; - // Add vectors one by one - builder copies directly to device memory for (float[] vector : data) { builder.addVector(vector); } From 8a6c6edbd8f6226b29f51e92476a8ef45eadac9d Mon Sep 17 00:00:00 2001 From: Vivek Narang Date: Fri, 22 May 2026 13:47:09 -0400 Subject: [PATCH 8/8] Adjust peak memory estimation for both NN_DESCENT and IVF_PQ --- ...ResourcesManager.java => CuVSResourcesManager.java} | 10 +++++----- .../lucene/Lucene99AcceleratedHNSWVectorsFormat.java | 4 ++-- .../lucene/Lucene99AcceleratedHNSWVectorsWriter.java | 6 +++--- ...eneAcceleratedHNSWBinaryQuantizedVectorsFormat.java | 4 ++-- ...eneAcceleratedHNSWBinaryQuantizedVectorsWriter.java | 6 +++--- ...eneAcceleratedHNSWScalarQuantizedVectorsFormat.java | 4 ++-- ...eneAcceleratedHNSWScalarQuantizedVectorsWriter.java | 6 +++--- 7 files changed, 20 insertions(+), 20 deletions(-) rename src/main/java/com/nvidia/cuvs/lucene/{CuvsResourcesManager.java => CuVSResourcesManager.java} (97%) diff --git a/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java b/src/main/java/com/nvidia/cuvs/lucene/CuVSResourcesManager.java similarity index 97% rename from src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java rename to src/main/java/com/nvidia/cuvs/lucene/CuVSResourcesManager.java index fc8d822c..38b8ef06 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/CuvsResourcesManager.java +++ b/src/main/java/com/nvidia/cuvs/lucene/CuVSResourcesManager.java @@ -29,12 +29,13 @@ * Manages a pool of finite {@link ManagedCuVSResources} and allows for the accessing threads * to lock and acquire available instance and release them back to the pool when finished. */ -public class CuvsResourcesManager { +public class CuVSResourcesManager { private static final Logger LOG = Logger.getLogger(Utils.class.getName()); private static final CuVSProvider PROVIDER = CuVSProvider.provider(); private static final GPUInfoProvider GPU_INFO_PROVIDER = PROVIDER.gpuInfoProvider(); private static final int MAX_POOL_SIZE = 512; + private static final double ONE_G = Math.pow(1024, 3); private ManagedCuVSResources[] pool; private ReentrantLock lock; @@ -44,7 +45,7 @@ public class CuvsResourcesManager { private int capacity; private CuVSResources cuVSResources; - public CuvsResourcesManager(int capacity) { + public CuVSResourcesManager(int capacity) { if (capacity > MAX_POOL_SIZE || capacity <= 0) { throw new IllegalArgumentException( "Invalid capacity, should be between 1 and " + MAX_POOL_SIZE); @@ -191,10 +192,9 @@ private long estimateNNDescentIndexBuildPeakMemory( */ final int sIdx = 4; - long datasetSize = rows * dimension * Float.BYTES; long nnDevicePeak = rows * (dimension * 2 + 276); long optimizePeak = rows * (4 + (sIdx + 1) * params.getIntermediateGraphDegree()); - return (long) ((datasetSize + Math.max(nnDevicePeak, optimizePeak)) * 1.25); + return (long) ((Math.max(nnDevicePeak, optimizePeak)) + ONE_G); } private long estimateIVFPQIndexBuildPeakMemory( @@ -246,7 +246,7 @@ private long estimateIVFPQIndexBuildPeakMemory( (LongStream.of((long) Math.ceil(ivfPQBuildPeak), ivfPQSearchPeak, optimizePeak) .max() .getAsLong() - * 1.25); + + (2.5 * ONE_G)); } /** diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java index df5e2941..ff7dd0c3 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsFormat.java @@ -31,7 +31,7 @@ public class Lucene99AcceleratedHNSWVectorsFormat extends KnnVectorsFormat { private static final FlatVectorsFormat FLAT_VECTORS_FORMAT; private static final int MAX_DIMENSIONS = 4096; private final AcceleratedHNSWParams acceleratedHNSWParams; - private final CuvsResourcesManager cuvsResourcesManager; + private final CuVSResourcesManager cuvsResourcesManager; static final String HNSW_META_CODEC_NAME = "Lucene99HnswVectorsFormatMeta"; static final String HNSW_META_CODEC_EXT = "vem"; @@ -69,7 +69,7 @@ public Lucene99AcceleratedHNSWVectorsFormat(AcceleratedHNSWParams acceleratedHNS super("Lucene99AcceleratedHNSWVectorsFormat"); this.acceleratedHNSWParams = acceleratedHNSWParams; if (isSupported()) { - cuvsResourcesManager = new CuvsResourcesManager(acceleratedHNSWParams.getWriterThreads()); + cuvsResourcesManager = new CuVSResourcesManager(acceleratedHNSWParams.getWriterThreads()); } else { cuvsResourcesManager = null; // Will not be needed in fallback mode. } diff --git a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 9465c311..0eb68ab9 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -23,7 +23,7 @@ import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; -import com.nvidia.cuvs.lucene.CuvsResourcesManager.ManagedCuVSResources; +import com.nvidia.cuvs.lucene.CuVSResourcesManager.ManagedCuVSResources; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -61,7 +61,7 @@ public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; - private final CuvsResourcesManager cuvsResourcesManager; + private final CuVSResourcesManager cuvsResourcesManager; private IndexOutput hnswMeta = null; private IndexOutput hnswVectorIndex = null; private String vemFileName; @@ -89,7 +89,7 @@ public Lucene99AcceleratedHNSWVectorsWriter( SegmentWriteState state, AcceleratedHNSWParams acceleratedHNSWParams, FlatVectorsWriter flatVectorsWriter, - CuvsResourcesManager cuvsResourcesManager) + CuVSResourcesManager cuvsResourcesManager) throws IOException { super(); this.flatVectorsWriter = flatVectorsWriter; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java index de2ce78a..efc5aebe 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat.java @@ -33,7 +33,7 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat extends KnnVector private static final int MAX_DIMENSIONS = 4096; private final AcceleratedHNSWParams acceleratedHNSWParams; - private final CuvsResourcesManager cuvsResourcesManager; + private final CuVSResourcesManager cuvsResourcesManager; static { try { @@ -65,7 +65,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat( super("Lucene99AcceleratedHNSWBinaryQuantizedVectorsFormat"); this.acceleratedHNSWParams = acceleratedHNSWParams; if (isSupported()) { - cuvsResourcesManager = new CuvsResourcesManager(acceleratedHNSWParams.getWriterThreads()); + cuvsResourcesManager = new CuVSResourcesManager(acceleratedHNSWParams.getWriterThreads()); } else { cuvsResourcesManager = null; // Will not be needed in fallback mode. } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 05e9bb97..467b5284 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -23,7 +23,7 @@ import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; -import com.nvidia.cuvs.lucene.CuvsResourcesManager.ManagedCuVSResources; +import com.nvidia.cuvs.lucene.CuVSResourcesManager.ManagedCuVSResources; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -63,7 +63,7 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter extends KnnVector private final List fields = new ArrayList<>(); private final InfoStream infoStream; private final AcceleratedHNSWParams acceleratedHNSWParams; - private final CuvsResourcesManager cuvsResourcesManager; + private final CuVSResourcesManager cuvsResourcesManager; private IndexOutput hnswMeta = null, hnswVectorIndex = null; private boolean finished; private String vemFileName; @@ -81,7 +81,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter( SegmentWriteState state, AcceleratedHNSWParams acceleratedHNSWParams, FlatVectorsWriter flatVectorsWriter, - CuvsResourcesManager cuvsResourcesManager) + CuVSResourcesManager cuvsResourcesManager) throws IOException { super(); this.acceleratedHNSWParams = acceleratedHNSWParams; diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java index 7cfb083a..383504fa 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsFormat.java @@ -30,7 +30,7 @@ public class LuceneAcceleratedHNSWScalarQuantizedVectorsFormat extends KnnVector private static final int MAX_DIMENSIONS = 4096; private final AcceleratedHNSWParams acceleratedHNSWParams; - private final CuvsResourcesManager cuvsResourcesManager; + private final CuVSResourcesManager cuvsResourcesManager; static { try { @@ -60,7 +60,7 @@ public LuceneAcceleratedHNSWScalarQuantizedVectorsFormat( super("Lucene99AcceleratedHNSWScalarQuantizedVectorsFormat"); this.acceleratedHNSWParams = acceleratedHNSWParams; if (isSupported()) { - cuvsResourcesManager = new CuvsResourcesManager(acceleratedHNSWParams.getWriterThreads()); + cuvsResourcesManager = new CuVSResourcesManager(acceleratedHNSWParams.getWriterThreads()); } else { cuvsResourcesManager = null; // Will not be needed in fallback mode. } diff --git a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 00875c6b..39d65ae3 100644 --- a/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -23,7 +23,7 @@ import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.lucene.AcceleratedHNSWUtils.QuantizationType; -import com.nvidia.cuvs.lucene.CuvsResourcesManager.ManagedCuVSResources; +import com.nvidia.cuvs.lucene.CuVSResourcesManager.ManagedCuVSResources; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -64,7 +64,7 @@ public class LuceneAcceleratedHNSWScalarQuantizedVectorsWriter extends KnnVector private final List fields = new ArrayList<>(); private final InfoStream infoStream; private final AcceleratedHNSWParams acceleratedHNSWParams; - private final CuvsResourcesManager cuvsResourcesManager; + private final CuVSResourcesManager cuvsResourcesManager; private IndexOutput hnswMeta = null, hnswVectorIndex = null; private boolean finished; private String vemFileName; @@ -91,7 +91,7 @@ public LuceneAcceleratedHNSWScalarQuantizedVectorsWriter( SegmentWriteState state, AcceleratedHNSWParams acceleratedHNSWParams, FlatVectorsWriter flatVectorsWriter, - CuvsResourcesManager cuvsResourcesManager) + CuVSResourcesManager cuvsResourcesManager) throws IOException { super(); this.acceleratedHNSWParams = acceleratedHNSWParams;