From 533ce9cf5eb42af12c5dcaf10d4cdb7d92d3e968 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lorenzo=20Dematt=C3=A9?= Date: Sat, 19 Jul 2025 06:46:46 +0200 Subject: [PATCH] [Java] Introduce scoped resource access (#1089) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR wraps access to the cuVS resources handle into a autocloseable accessor. The reason is to provide a clear scope for resource access, with the possibility to add application logic to the scope begin/end. An example of this logic is the `CheckedCuVSResources` decorator introduced in the test suite, which checks that the `CuVSResources` it decorates has not been destroyed yet, and that access to the underlying cuVS resources handle is not concurrent (is not done by more than 1 thread at every time). I preferred to write it as a decorator, as this gives the API consumers more freedom: if they want the check, they can use `CheckedCuVSResources` (or a similar class); if they want direct, straight, fast access, they can just use the accessor; if they need more logic (e.g. logging or anything else), that is possible too. Tests are updated to use `CheckedCuVSResources`, so we can spot issues quickly in tests without burdening production code with additional logic. Authors: - Lorenzo Dematté (https://github.com/ldematte) - MithunR (https://github.com/mythrocks) Approvers: - MithunR (https://github.com/mythrocks) URL: https://github.com/rapidsai/cuvs/pull/1089 --- .../java/com/nvidia/cuvs/CuVSResources.java | 26 +- .../cuvs/internal/BruteForceIndexImpl.java | 238 +++++----- .../nvidia/cuvs/internal/CagraIndexImpl.java | 255 +++++----- .../cuvs/internal/CuVSResourcesImpl.java | 37 +- .../nvidia/cuvs/internal/HnswIndexImpl.java | 74 ++- .../nvidia/cuvs/internal/TieredIndexImpl.java | 439 ++++++++---------- .../nvidia/cuvs/BruteForceAndSearchIT.java | 8 +- .../nvidia/cuvs/BruteForceRandomizedIT.java | 2 +- .../nvidia/cuvs/CagraBuildAndSearchIT.java | 24 +- .../com/nvidia/cuvs/CagraRandomizedIT.java | 2 +- .../com/nvidia/cuvs/CheckedCuVSResources.java | 74 +++ .../java/com/nvidia/cuvs/CuVSResourcesIT.java | 98 ++++ .../com/nvidia/cuvs/HnswBuildAndSearchIT.java | 6 +- .../com/nvidia/cuvs/HnswRandomizedIT.java | 2 +- 14 files changed, 719 insertions(+), 566 deletions(-) create mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java create mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSResourcesIT.java diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.java index 6a84f737b1..2d362b6c4c 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSResources.java @@ -26,7 +26,31 @@ public interface CuVSResources extends AutoCloseable { /** - * Closes this resources and releases any resources associated with it. + * Provide scoped access to the native resources object. + */ + interface ScopedAccess extends AutoCloseable { + /** + * Gets the opaque CuVSResources handle, to be used whenever we need to pass a cuvsResources_t parameter + * + * @return the CuVSResources handle + */ + long handle(); + + @Override + void close(); + } + + /** + * Gets scoped access to the native resources object. + * The native resource object is not thread safe: only a single thread at every time should access + * concurrently the same native resources. Calling this method from multiple thread is OK, but the + * returned {@link ScopedAccess} object must be closed before calling {@code access()} again from a + * different thread. + */ + ScopedAccess access(); + + /** + * Closes this CuVSResources object and releases any resources associated with it. */ @Override void close(); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java index f9e3b9ba93..163205b830 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/BruteForceIndexImpl.java @@ -67,7 +67,7 @@ */ public class BruteForceIndexImpl implements BruteForceIndex { - private final CuVSResourcesImpl resources; + private final CuVSResources resources; private final IndexReference bruteForceIndexReference; private boolean destroyed; @@ -76,12 +76,12 @@ public class BruteForceIndexImpl implements BruteForceIndex { * * @param dataset the dataset used for creating the BRUTEFORCE * index - * @param resources an instance of {@link CuVSResourcesImpl} + * @param resources an instance of {@link CuVSResources} * @param bruteForceIndexParams an instance of {@link BruteForceIndexParams} * holding the index parameters */ private BruteForceIndexImpl( - Dataset dataset, CuVSResourcesImpl resources, BruteForceIndexParams bruteForceIndexParams) + Dataset dataset, CuVSResources resources, BruteForceIndexParams bruteForceIndexParams) throws Exception { Objects.requireNonNull(dataset); try (dataset) { @@ -95,10 +95,9 @@ private BruteForceIndexImpl( * Constructor for loading the index from an {@link InputStream} * * @param inputStream an instance of stream to read the index bytes from - * @param resources an instance of {@link CuVSResourcesImpl} + * @param resources an instance of {@link CuVSResources} */ - private BruteForceIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) - throws Throwable { + private BruteForceIndexImpl(InputStream inputStream, CuVSResources resources) throws Throwable { this.resources = resources; this.bruteForceIndexReference = deserialize(inputStream); } @@ -121,12 +120,14 @@ public void destroyIndex() { checkCuVSError(returnValue, "cuvsBruteForceIndexDestroy"); if (bruteForceIndexReference.datasetBytes > 0) { - returnValue = - cuvsRMMFree( - resources.getHandle(), - bruteForceIndexReference.datasetPtr, - bruteForceIndexReference.datasetBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); + try (var resourcesAccessor = resources.access()) { + checkCuVSError( + cuvsRMMFree( + resourcesAccessor.handle(), + bruteForceIndexReference.datasetPtr, + bruteForceIndexReference.datasetBytes), + "cuvsRMMFree"); + } } if (bruteForceIndexReference.tensorDataArena != null) { bruteForceIndexReference.tensorDataArena.close(); @@ -141,7 +142,7 @@ public void destroyIndex() { * build the {@link BruteForceIndex} * * @return an instance of {@link IndexReference} that holds the pointer to the - * index + * index */ private IndexReference build(DatasetImpl dataset, BruteForceIndexParams bruteForceIndexParams) { long rows = dataset.size(); @@ -149,36 +150,35 @@ private IndexReference build(DatasetImpl dataset, BruteForceIndexParams bruteFor MemorySegment datasetMemSegment = dataset.asMemorySegment(); - long cuvsResources = resources.getHandle(); - omp_set_num_threads(bruteForceIndexParams.getNumWriterThreads()); - long datasetBytes = C_FLOAT_BYTE_SIZE * rows * cols; - MemorySegment datasetMemorySegmentP = allocateRMMSegment(cuvsResources, datasetBytes); - cudaMemcpy(datasetMemorySegmentP, datasetMemSegment, datasetBytes, INFER_DIRECTION); + long datasetBytes = C_FLOAT_BYTE_SIZE * rows * cols; + var index = createBruteForceIndex(); - long[] datasetShape = {rows, cols}; - var tensorDataArena = Arena.ofShared(); - MemorySegment datasetTensor = - prepareTensor(tensorDataArena, datasetMemorySegmentP, datasetShape, 2, 32, 2, 2, 1); + try (var resourcesAccessor = resources.access()) { + long cuvsResources = resourcesAccessor.handle(); + MemorySegment datasetMemorySegmentP = allocateRMMSegment(cuvsResources, datasetBytes); - var indexReference = - new IndexReference( - datasetMemorySegmentP, datasetBytes, tensorDataArena, createBruteForceIndex()); + cudaMemcpy(datasetMemorySegmentP, datasetMemSegment, datasetBytes, INFER_DIRECTION); - var returnValue = cuvsStreamSync(cuvsResources); - checkCuVSError(returnValue, "cuvsStreamSync"); + long[] datasetShape = {rows, cols}; + var tensorDataArena = Arena.ofShared(); + MemorySegment datasetTensor = + prepareTensor(tensorDataArena, datasetMemorySegmentP, datasetShape, 2, 32, 2, 2, 1); - returnValue = - cuvsBruteForceBuild(cuvsResources, datasetTensor, 0, 0.0f, indexReference.indexPtr); - checkCuVSError(returnValue, "cuvsBruteForceBuild"); + var returnValue = cuvsStreamSync(cuvsResources); + checkCuVSError(returnValue, "cuvsStreamSync"); - returnValue = cuvsStreamSync(cuvsResources); - checkCuVSError(returnValue, "cuvsStreamSync"); + returnValue = cuvsBruteForceBuild(cuvsResources, datasetTensor, 0, 0.0f, index); + checkCuVSError(returnValue, "cuvsBruteForceBuild"); - omp_set_num_threads(1); + returnValue = cuvsStreamSync(cuvsResources); + checkCuVSError(returnValue, "cuvsStreamSync"); - return indexReference; + return new IndexReference(datasetMemorySegmentP, datasetBytes, tensorDataArena, index); + } finally { + omp_set_num_threads(1); + } } /** @@ -203,94 +203,98 @@ public SearchResults search(BruteForceQuery cuvsQuery) throws Throwable { MemorySegment distancesMemorySegment = localArena.allocate(distancesSequenceLayout); // prepare the prefiltering data - long prefilterDataLength = 0; - MemorySegment prefilterDataMemorySegment = MemorySegment.NULL; + final long prefilterDataLength; + final MemorySegment prefilterDataMemorySegment; BitSet[] prefilters = cuvsQuery.getPrefilters(); if (prefilters != null && prefilters.length > 0) { BitSet concatenatedFilters = concatenate(prefilters, cuvsQuery.getNumDocs()); long[] filters = concatenatedFilters.toLongArray(); prefilterDataMemorySegment = buildMemorySegment(localArena, filters); prefilterDataLength = (long) cuvsQuery.getNumDocs() * prefilters.length; + } else { + prefilterDataLength = 0; + prefilterDataMemorySegment = MemorySegment.NULL; } MemorySegment querySeg = buildMemorySegment(localArena, cuvsQuery.getQueryVectors()); int topk = cuvsQuery.getTopK(); - long cuvsResources = resources.getHandle(); - - long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; - long neighborsBytes = C_LONG_BYTE_SIZE * numQueries * topk; - long distanceBytes = C_FLOAT_BYTE_SIZE * numQueries * topk; - long prefilterBytes = 0; // size assigned later - - MemorySegment queriesDP = allocateRMMSegment(cuvsResources, queriesBytes); - MemorySegment neighborsDP = allocateRMMSegment(cuvsResources, neighborsBytes); - MemorySegment distancesDP = allocateRMMSegment(cuvsResources, distanceBytes); - MemorySegment prefilterDP = MemorySegment.NULL; - - cudaMemcpy(queriesDP, querySeg, queriesBytes, INFER_DIRECTION); - - long[] queriesShape = {numQueries, vectorDimension}; - MemorySegment queriesTensor = - prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 2, 1); - long[] neighborsShape = {numQueries, topk}; - MemorySegment neighborsTensor = - prepareTensor(localArena, neighborsDP, neighborsShape, 0, 64, 2, 2, 1); - long[] distancesShape = {numQueries, topk}; - MemorySegment distancesTensor = - prepareTensor(localArena, distancesDP, distancesShape, 2, 32, 2, 2, 1); - - MemorySegment prefilter = cuvsFilter.allocate(localArena); - MemorySegment prefilterTensor; - - if (prefilterDataMemorySegment == MemorySegment.NULL) { - cuvsFilter.type(prefilter, 0); // NO_FILTER - cuvsFilter.addr(prefilter, 0); - } else { - long[] prefilterShape = {(prefilterDataLength + 31) / 32}; - long prefilterLen = prefilterShape[0]; - prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; + try (var resourcesAccessor = resources.access()) { + long cuvsResources = resourcesAccessor.handle(); - prefilterDP = allocateRMMSegment(cuvsResources, prefilterBytes); + long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; + long neighborsBytes = C_LONG_BYTE_SIZE * numQueries * topk; + long distanceBytes = C_FLOAT_BYTE_SIZE * numQueries * topk; + long prefilterBytes = 0; // size assigned later - cudaMemcpy(prefilterDP, prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE); + MemorySegment queriesDP = allocateRMMSegment(cuvsResources, queriesBytes); + MemorySegment neighborsDP = allocateRMMSegment(cuvsResources, neighborsBytes); + MemorySegment distancesDP = allocateRMMSegment(cuvsResources, distanceBytes); + MemorySegment prefilterDP = MemorySegment.NULL; - prefilterTensor = prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 1, 2, 1); + cudaMemcpy(queriesDP, querySeg, queriesBytes, INFER_DIRECTION); - cuvsFilter.type(prefilter, 2); - cuvsFilter.addr(prefilter, prefilterTensor.address()); - } + long[] queriesShape = {numQueries, vectorDimension}; + MemorySegment queriesTensor = + prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 2, 1); + long[] neighborsShape = {numQueries, topk}; + MemorySegment neighborsTensor = + prepareTensor(localArena, neighborsDP, neighborsShape, 0, 64, 2, 2, 1); + long[] distancesShape = {numQueries, topk}; + MemorySegment distancesTensor = + prepareTensor(localArena, distancesDP, distancesShape, 2, 32, 2, 2, 1); - var returnValue = cuvsStreamSync(cuvsResources); - checkCuVSError(returnValue, "cuvsStreamSync"); + MemorySegment prefilter = cuvsFilter.allocate(localArena); + MemorySegment prefilterTensor; - returnValue = - cuvsBruteForceSearch( - cuvsResources, - bruteForceIndexReference.indexPtr, - queriesTensor, - neighborsTensor, - distancesTensor, - prefilter); - checkCuVSError(returnValue, "cuvsBruteForceSearch"); + if (prefilterDataMemorySegment == MemorySegment.NULL) { + cuvsFilter.type(prefilter, 0); // NO_FILTER + cuvsFilter.addr(prefilter, 0); + } else { + long[] prefilterShape = {(prefilterDataLength + 31) / 32}; + long prefilterLen = prefilterShape[0]; + prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; - returnValue = cuvsStreamSync(cuvsResources); - checkCuVSError(returnValue, "cuvsStreamSync"); + prefilterDP = allocateRMMSegment(cuvsResources, prefilterBytes); + + cudaMemcpy(prefilterDP, prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE); + + prefilterTensor = prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 1, 2, 1); + + cuvsFilter.type(prefilter, 2); + cuvsFilter.addr(prefilter, prefilterTensor.address()); + } - cudaMemcpy(neighborsMemorySegment, neighborsDP, neighborsBytes, INFER_DIRECTION); - cudaMemcpy(distancesMemorySegment, distancesDP, distanceBytes, INFER_DIRECTION); - - returnValue = cuvsRMMFree(cuvsResources, neighborsDP, neighborsBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - returnValue = cuvsRMMFree(cuvsResources, distancesDP, distanceBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - returnValue = cuvsRMMFree(cuvsResources, queriesDP, queriesBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - if (prefilterBytes > 0) { - returnValue = cuvsRMMFree(cuvsResources, prefilterDP, prefilterBytes); + var returnValue = cuvsStreamSync(cuvsResources); + checkCuVSError(returnValue, "cuvsStreamSync"); + + returnValue = + cuvsBruteForceSearch( + cuvsResources, + bruteForceIndexReference.indexPtr, + queriesTensor, + neighborsTensor, + distancesTensor, + prefilter); + checkCuVSError(returnValue, "cuvsBruteForceSearch"); + + returnValue = cuvsStreamSync(cuvsResources); + checkCuVSError(returnValue, "cuvsStreamSync"); + + cudaMemcpy(neighborsMemorySegment, neighborsDP, neighborsBytes, INFER_DIRECTION); + cudaMemcpy(distancesMemorySegment, distancesDP, distanceBytes, INFER_DIRECTION); + + returnValue = cuvsRMMFree(cuvsResources, neighborsDP, neighborsBytes); + checkCuVSError(returnValue, "cuvsRMMFree"); + returnValue = cuvsRMMFree(cuvsResources, distancesDP, distanceBytes); + checkCuVSError(returnValue, "cuvsRMMFree"); + returnValue = cuvsRMMFree(cuvsResources, queriesDP, queriesBytes); checkCuVSError(returnValue, "cuvsRMMFree"); + if (prefilterBytes > 0) { + returnValue = cuvsRMMFree(cuvsResources, prefilterDP, prefilterBytes); + checkCuVSError(returnValue, "cuvsRMMFree"); + } } - return BruteForceSearchResults.create( neighborsSequenceLayout, distancesSequenceLayout, @@ -312,19 +316,19 @@ public void serialize(OutputStream outputStream) throws Throwable { @Override public void serialize(OutputStream outputStream, Path tempFile) throws Throwable { checkNotDestroyed(); - tempFile = tempFile.toAbsolutePath(); + var tempFilePath = tempFile.toAbsolutePath(); - long cuvsRes = resources.getHandle(); - try (var localArena = Arena.ofConfined()) { + try (var localArena = Arena.ofConfined(); + var resourcesAccessor = resources.access()) { int returnValue = cuvsBruteForceSerialize( - cuvsRes, - localArena.allocateFrom(tempFile.toString()), + resourcesAccessor.handle(), + localArena.allocateFrom(tempFilePath.toString()), bruteForceIndexReference.indexPtr); checkCuVSError(returnValue, "cuvsBruteForceSerialize"); } - try (var inputStream = Files.newInputStream(tempFile)) { + try (var inputStream = Files.newInputStream(tempFilePath)) { inputStream.transferTo(outputStream); } finally { Files.deleteIfExists(tempFile); @@ -363,13 +367,15 @@ private IndexReference deserialize(InputStream inputStream) throws Throwable { try (inputStream; var outputStream = Files.newOutputStream(tmpIndexFile); - var arena = Arena.ofConfined()) { + var arena = Arena.ofConfined(); + var resourcesAccessor = resources.access()) { inputStream.transferTo(outputStream); - long cuvsRes = resources.getHandle(); int returnValue = cuvsBruteForceDeserialize( - cuvsRes, arena.allocateFrom(tmpIndexFile.toString()), indexReference.indexPtr); + resourcesAccessor.handle(), + arena.allocateFrom(tmpIndexFile.toString()), + indexReference.indexPtr); checkCuVSError(returnValue, "cuvsBruteForceDeserialize"); } finally { @@ -379,11 +385,7 @@ private IndexReference deserialize(InputStream inputStream) throws Throwable { } public static BruteForceIndex.Builder newBuilder(CuVSResources cuvsResources) { - Objects.requireNonNull(cuvsResources); - if (!(cuvsResources instanceof CuVSResourcesImpl)) { - throw new IllegalArgumentException("Unsupported " + cuvsResources); - } - return new Builder((CuVSResourcesImpl) cuvsResources); + return new Builder(Objects.requireNonNull(cuvsResources)); } /** @@ -392,16 +394,16 @@ public static BruteForceIndex.Builder newBuilder(CuVSResources cuvsResources) { public static class Builder implements BruteForceIndex.Builder { private Dataset dataset; - private final CuVSResourcesImpl cuvsResources; + private final CuVSResources cuvsResources; private BruteForceIndexParams bruteForceIndexParams; private InputStream inputStream; /** - * Constructs this Builder with an instance of {@link CuVSResourcesImpl}. + * Constructs this Builder with an instance of {@link CuVSResources}. * * @param cuvsResources an instance of {@link CuVSResources} */ - public Builder(CuVSResourcesImpl cuvsResources) { + public Builder(CuVSResources cuvsResources) { this.cuvsResources = cuvsResources; } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 49be7104e8..0e1fcb9048 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -87,7 +87,7 @@ * @since 25.02 */ public class CagraIndexImpl implements CagraIndex { - private final CuVSResourcesImpl resources; + private final CuVSResources resources; private final IndexReference cagraIndexReference; private boolean destroyed; @@ -100,7 +100,7 @@ public class CagraIndexImpl implements CagraIndex { * @param resources an instance of {@link CuVSResources} */ private CagraIndexImpl( - CagraIndexParams indexParameters, Dataset dataset, CuVSResourcesImpl resources) { + CagraIndexParams indexParameters, Dataset dataset, CuVSResources resources) { Objects.requireNonNull(dataset); this.resources = resources; assert dataset instanceof DatasetImpl; @@ -113,7 +113,7 @@ private CagraIndexImpl( * @param inputStream an instance of stream to read the index bytes from * @param resources an instance of {@link CuVSResources} */ - private CagraIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) throws Throwable { + private CagraIndexImpl(InputStream inputStream, CuVSResources resources) throws Throwable { this.resources = resources; this.cagraIndexReference = deserialize(inputStream); } @@ -125,7 +125,7 @@ private CagraIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) thr * @param indexReference The reference to the existing index * @param resources The resources instance */ - private CagraIndexImpl(IndexReference indexReference, CuVSResourcesImpl resources) { + private CagraIndexImpl(IndexReference indexReference, CuVSResources resources) { this.resources = resources; this.cagraIndexReference = indexReference; this.destroyed = false; @@ -176,8 +176,6 @@ private IndexReference build(CagraIndexParams indexParameters, DatasetImpl datas MemorySegment dataSeg = dataset.asMemorySegment(); - long cuvsRes = resources.getHandle(); - long[] datasetShape = {rows, cols}; MemorySegment datasetTensor = prepareTensor(localArena, dataSeg, datasetShape, 2, 32, 2, 2, 1); @@ -194,16 +192,18 @@ private IndexReference build(CagraIndexParams indexParameters, DatasetImpl datas cuvsIvfPqIndexParams.n_lists( cuvsIvfPqIndexParamsMS, (int) (rows < n_lists ? rows : n_lists)); } + try (var resourcesAccessor = resources.access()) { + var cuvsRes = resourcesAccessor.handle(); - var returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); - - returnValue = cuvsCagraBuild(cuvsRes, indexParamsMemorySegment, datasetTensor, index); - checkCuVSError(returnValue, "cuvsCagraBuild"); + var returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); + returnValue = cuvsCagraBuild(cuvsRes, indexParamsMemorySegment, datasetTensor, index); + checkCuVSError(returnValue, "cuvsCagraBuild"); + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + } omp_set_num_threads(1); return new IndexReference(index, dataset); @@ -248,97 +248,101 @@ public SearchResults search(CagraQuery query) throws Throwable { MemorySegment distancesMemorySegment = localArena.allocate(distancesSequenceLayout); MemorySegment floatsSeg = buildMemorySegment(localArena, query.getQueryVectors()); - long cuvsRes = resources.getHandle(); - long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; long neighborsBytes = C_INT_BYTE_SIZE * numQueries * topK; long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; - long prefilterBytes = 0; // size assigned later - - MemorySegment queriesDP = allocateRMMSegment(cuvsRes, queriesBytes); - MemorySegment neighborsDP = allocateRMMSegment(cuvsRes, neighborsBytes); - MemorySegment distancesDP = allocateRMMSegment(cuvsRes, distancesBytes); - MemorySegment prefilterDP = MemorySegment.NULL; - long prefilterLen = 0; - - cudaMemcpy(queriesDP, floatsSeg, queriesBytes, INFER_DIRECTION); - - long[] queriesShape = {numQueries, vectorDimension}; - MemorySegment queriesTensor = - prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 2, 1); - long[] neighborsShape = {numQueries, topK}; - MemorySegment neighborsTensor = - prepareTensor(localArena, neighborsDP, neighborsShape, 1, 32, 2, 2, 1); - long[] distancesShape = {numQueries, topK}; - MemorySegment distancesTensor = - prepareTensor(localArena, distancesDP, distancesShape, 2, 32, 2, 2, 1); - - var returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); - - // prepare the prefiltering data - long prefilterDataLength = 0; - MemorySegment prefilterDataMemorySegment = MemorySegment.NULL; - BitSet[] prefilters; - if (query.getPrefilter() != null) { - prefilters = new BitSet[] {query.getPrefilter()}; - BitSet concatenatedFilters = concatenate(prefilters, query.getNumDocs()); - long[] filters = concatenatedFilters.toLongArray(); - prefilterDataMemorySegment = buildMemorySegment(localArena, filters); - prefilterDataLength = query.getNumDocs() * prefilters.length; - } - MemorySegment prefilter = cuvsFilter.allocate(localArena); - MemorySegment prefilterTensor; + try (var resourcesAccessor = resources.access()) { + var cuvsRes = resourcesAccessor.handle(); + + MemorySegment queriesDP = allocateRMMSegment(cuvsRes, queriesBytes); + MemorySegment neighborsDP = allocateRMMSegment(cuvsRes, neighborsBytes); + MemorySegment distancesDP = allocateRMMSegment(cuvsRes, distancesBytes); + MemorySegment prefilterDP = MemorySegment.NULL; + long prefilterLen = 0; + + cudaMemcpy(queriesDP, floatsSeg, queriesBytes, INFER_DIRECTION); + + long[] queriesShape = {numQueries, vectorDimension}; + MemorySegment queriesTensor = + prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 2, 1); + long[] neighborsShape = {numQueries, topK}; + MemorySegment neighborsTensor = + prepareTensor(localArena, neighborsDP, neighborsShape, 1, 32, 2, 2, 1); + long[] distancesShape = {numQueries, topK}; + MemorySegment distancesTensor = + prepareTensor(localArena, distancesDP, distancesShape, 2, 32, 2, 2, 1); + + var returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + // prepare the prefiltering data + long prefilterDataLength = 0; + MemorySegment prefilterDataMemorySegment = MemorySegment.NULL; + BitSet[] prefilters; + if (query.getPrefilter() != null) { + prefilters = new BitSet[] {query.getPrefilter()}; + BitSet concatenatedFilters = concatenate(prefilters, query.getNumDocs()); + long[] filters = concatenatedFilters.toLongArray(); + prefilterDataMemorySegment = buildMemorySegment(localArena, filters); + prefilterDataLength = query.getNumDocs() * prefilters.length; + } - if (prefilterDataMemorySegment == MemorySegment.NULL) { - cuvsFilter.type(prefilter, 0); // NO_FILTER - cuvsFilter.addr(prefilter, 0); - } else { - long[] prefilterShape = {(prefilterDataLength + 31) / 32}; - prefilterLen = prefilterShape[0]; - prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; + MemorySegment prefilter = cuvsFilter.allocate(localArena); + MemorySegment prefilterTensor; - prefilterDP = allocateRMMSegment(cuvsRes, prefilterBytes); + final long prefilterBytes; - cudaMemcpy(prefilterDP, prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE); + if (prefilterDataMemorySegment == MemorySegment.NULL) { + cuvsFilter.type(prefilter, 0); // NO_FILTER + cuvsFilter.addr(prefilter, 0); + prefilterBytes = 0; + } else { + long[] prefilterShape = {(prefilterDataLength + 31) / 32}; + prefilterLen = prefilterShape[0]; + prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; - prefilterTensor = prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 1, 2, 1); + prefilterDP = allocateRMMSegment(cuvsRes, prefilterBytes); - cuvsFilter.type(prefilter, 1); - cuvsFilter.addr(prefilter, prefilterTensor.address()); - } + cudaMemcpy(prefilterDP, prefilterDataMemorySegment, prefilterBytes, HOST_TO_DEVICE); - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); + prefilterTensor = prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 1, 2, 1); - returnValue = - cuvsCagraSearch( - cuvsRes, - segmentFromSearchParams(localArena, query.getCagraSearchParameters()), - cagraIndexReference.getMemorySegment(), - queriesTensor, - neighborsTensor, - distancesTensor, - prefilter); - checkCuVSError(returnValue, "cuvsCagraSearch"); - - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); - - cudaMemcpy(neighborsMemorySegment, neighborsDP, neighborsBytes, INFER_DIRECTION); - cudaMemcpy(distancesMemorySegment, distancesDP, distancesBytes, INFER_DIRECTION); - - returnValue = cuvsRMMFree(cuvsRes, distancesDP, distancesBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - returnValue = cuvsRMMFree(cuvsRes, neighborsDP, neighborsBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - returnValue = cuvsRMMFree(cuvsRes, queriesDP, queriesBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - - if (prefilterLen > 0) { - returnValue = cuvsRMMFree(cuvsRes, prefilterDP, C_INT_BYTE_SIZE * prefilterBytes); + cuvsFilter.type(prefilter, 1); + cuvsFilter.addr(prefilter, prefilterTensor.address()); + } + + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + returnValue = + cuvsCagraSearch( + cuvsRes, + segmentFromSearchParams(localArena, query.getCagraSearchParameters()), + cagraIndexReference.getMemorySegment(), + queriesTensor, + neighborsTensor, + distancesTensor, + prefilter); + checkCuVSError(returnValue, "cuvsCagraSearch"); + + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + cudaMemcpy(neighborsMemorySegment, neighborsDP, neighborsBytes, INFER_DIRECTION); + cudaMemcpy(distancesMemorySegment, distancesDP, distancesBytes, INFER_DIRECTION); + + returnValue = cuvsRMMFree(cuvsRes, distancesDP, distancesBytes); + checkCuVSError(returnValue, "cuvsRMMFree"); + returnValue = cuvsRMMFree(cuvsRes, neighborsDP, neighborsBytes); checkCuVSError(returnValue, "cuvsRMMFree"); + returnValue = cuvsRMMFree(cuvsRes, queriesDP, queriesBytes); + checkCuVSError(returnValue, "cuvsRMMFree"); + + if (prefilterLen > 0) { + returnValue = cuvsRMMFree(cuvsRes, prefilterDP, C_INT_BYTE_SIZE * prefilterBytes); + checkCuVSError(returnValue, "cuvsRMMFree"); + } } return CagraSearchResults.create( @@ -370,26 +374,27 @@ public void serialize(OutputStream outputStream, int bufferLength) throws Throwa public void serialize(OutputStream outputStream, Path tempFile, int bufferLength) throws Throwable { checkNotDestroyed(); - tempFile = tempFile.toAbsolutePath(); - try (var localArena = Arena.ofConfined()) { + final var tempFilePath = tempFile.toAbsolutePath(); + try (var localArena = Arena.ofConfined(); + var resourcesAccessor = resources.access()) { - long cuvsRes = resources.getHandle(); + long cuvsRes = resourcesAccessor.handle(); var returnValue = cuvsCagraSerialize( cuvsRes, - localArena.allocateFrom(tempFile.toString()), + localArena.allocateFrom(tempFilePath.toString()), cagraIndexReference.getMemorySegment(), true); checkCuVSError(returnValue, "cuvsCagraSerialize"); - try (var fileInputStream = Files.newInputStream(tempFile)) { + try (var fileInputStream = Files.newInputStream(tempFilePath)) { byte[] chunk = new byte[bufferLength]; int chunkLength = 0; while ((chunkLength = fileInputStream.read(chunk)) != -1) { outputStream.write(chunk, 0, chunkLength); } } finally { - Files.deleteIfExists(tempFile); + Files.deleteIfExists(tempFilePath); } } } @@ -412,25 +417,27 @@ public void serializeToHNSW(OutputStream outputStream, int bufferLength) throws public void serializeToHNSW(OutputStream outputStream, Path tempFile, int bufferLength) throws Throwable { checkNotDestroyed(); - tempFile = tempFile.toAbsolutePath(); + final var tempFilePath = tempFile.toAbsolutePath(); try (var localArena = Arena.ofConfined()) { MemorySegment pathSeg = buildMemorySegment(localArena, tempFile.toString()); - long cuvsRes = resources.getHandle(); - int returnValue = - cuvsCagraSerializeToHnswlib(cuvsRes, pathSeg, cagraIndexReference.getMemorySegment()); - checkCuVSError(returnValue, "cuvsCagraSerializeToHnswlib"); + try (var resourcesAccessor = resources.access()) { + checkCuVSError( + cuvsCagraSerializeToHnswlib( + resourcesAccessor.handle(), pathSeg, cagraIndexReference.getMemorySegment()), + "cuvsCagraSerializeToHnswlib"); + } } - try (FileInputStream fileInputStream = new FileInputStream(tempFile.toFile())) { + try (FileInputStream fileInputStream = new FileInputStream(tempFilePath.toFile())) { byte[] chunk = new byte[bufferLength]; int chunkLength; while ((chunkLength = fileInputStream.read(chunk)) != -1) { outputStream.write(chunk, 0, chunkLength); } } finally { - Files.deleteIfExists(tempFile); + Files.deleteIfExists(tempFilePath); } } @@ -452,11 +459,12 @@ private IndexReference deserialize(InputStream inputStream) throws Throwable { var arena = Arena.ofConfined()) { inputStream.transferTo(outputStream); - long cuvsRes = resources.getHandle(); - var returnValue = - cuvsCagraDeserialize(cuvsRes, arena.allocateFrom(tmpIndexFile.toString()), index); - checkCuVSError(returnValue, "cuvsCagraDeserialize"); - + try (var resourcesAccessor = resources.access()) { + var cuvsRes = resourcesAccessor.handle(); + var returnValue = + cuvsCagraDeserialize(cuvsRes, arena.allocateFrom(tmpIndexFile.toString()), index); + checkCuVSError(returnValue, "cuvsCagraDeserialize"); + } } finally { Files.deleteIfExists(tmpIndexFile); } @@ -469,7 +477,7 @@ private IndexReference deserialize(InputStream inputStream) throws Throwable { * @return an instance of {@link CuVSResources} */ @Override - public CuVSResourcesImpl getCuVSResources() { + public CuVSResources getCuVSResources() { return resources; } @@ -582,11 +590,7 @@ private MemorySegment segmentFromSearchParams(Arena arena, CagraSearchParams par } public static CagraIndex.Builder newBuilder(CuVSResources cuvsResources) { - Objects.requireNonNull(cuvsResources); - if (!(cuvsResources instanceof CuVSResourcesImpl)) { - throw new IllegalArgumentException("Unsupported " + cuvsResources); - } - return new CagraIndexImpl.Builder((CuVSResourcesImpl) cuvsResources); + return new CagraIndexImpl.Builder(Objects.requireNonNull(cuvsResources)); } /** @@ -607,9 +611,8 @@ public static CagraIndex merge(CagraIndex[] indexes) { * @return A new merged CAGRA index */ public static CagraIndex merge(CagraIndex[] indexes, CagraMergeParams mergeParams) { - CuVSResourcesImpl resources = (CuVSResourcesImpl) indexes[0].getCuVSResources(); + CuVSResources resources = indexes[0].getCuVSResources(); var mergedIndex = createCagraIndex(); - long cuvsRes = resources.getHandle(); try (var localArena = Arena.ofConfined()) { MemorySegment indexesSegment = @@ -624,11 +627,13 @@ public static CagraIndex merge(CagraIndex[] indexes, CagraMergeParams mergeParam // TODO: we should call cuvsCreateMergeParams here, instead of allocating this ourselves // See https://github.com/rapidsai/cuvs/pull/1109 var mergeParamsSegment = createMergeParamsSegment(localArena, mergeParams); - - var returnValue = - cuvsCagraMerge(cuvsRes, mergeParamsSegment, indexesSegment, indexes.length, mergedIndex); - - checkCuVSError(returnValue, "cuvsCagraMerge"); + try (var resourcesAccessor = resources.access()) { + var cuvsRes = resourcesAccessor.handle(); + var returnValue = + cuvsCagraMerge( + cuvsRes, mergeParamsSegment, indexesSegment, indexes.length, mergedIndex); + checkCuVSError(returnValue, "cuvsCagraMerge"); + } } return new CagraIndexImpl(new IndexReference(mergedIndex, null), resources); @@ -670,10 +675,10 @@ public static class Builder implements CagraIndex.Builder { private Dataset dataset; private CagraIndexParams cagraIndexParams; - private final CuVSResourcesImpl cuvsResources; + private final CuVSResources cuvsResources; private InputStream inputStream; - public Builder(CuVSResourcesImpl cuvsResources) { + public Builder(CuVSResources cuvsResources) { this.cuvsResources = cuvsResources; } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java index 5acb2cb9bd..8f9e1d94d0 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSResourcesImpl.java @@ -33,7 +33,7 @@ public class CuVSResourcesImpl implements CuVSResources { private final Path tempDirectory; private final long resourceHandle; - private boolean destroyed; + private final ScopedAccess access; /** * Constructor that allocates the resources needed for cuVS @@ -45,17 +45,30 @@ public CuVSResourcesImpl(Path tempDirectory) { var resourcesMemorySegment = localArena.allocate(cuvsResources_t); int returnValue = cuvsResourcesCreate(resourcesMemorySegment); checkCuVSError(returnValue, "cuvsResourcesCreate"); - resourceHandle = resourcesMemorySegment.get(cuvsResources_t, 0); + this.resourceHandle = resourcesMemorySegment.get(cuvsResources_t, 0); + this.access = + new ScopedAccess() { + @Override + public long handle() { + return resourceHandle; + } + + @Override + public void close() {} + }; } } + @Override + public ScopedAccess access() { + return this.access; + } + @Override public void close() { synchronized (this) { - checkNotDestroyed(); int returnValue = cuvsResourcesDestroy(resourceHandle); checkCuVSError(returnValue, "cuvsResourcesDestroy"); - destroyed = true; } } @@ -63,20 +76,4 @@ public void close() { public Path tempDirectory() { return tempDirectory; } - - private void checkNotDestroyed() { - if (destroyed) { - throw new IllegalStateException("destroyed"); - } - } - - /** - * Gets the opaque CuVSResources handle, to be used whenever we need to pass a cuvsResources_t parameter - * - * @return the CuVSResources handle - */ - long getHandle() { - checkNotDestroyed(); - return resourceHandle; - } } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java index 178ed21b87..0d8997dca8 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java @@ -55,7 +55,7 @@ */ public class HnswIndexImpl implements HnswIndex { - private final CuVSResourcesImpl resources; + private final CuVSResources resources; private final HnswIndexParams hnswIndexParams; private final IndexReference hnswIndexReference; @@ -63,10 +63,10 @@ public class HnswIndexImpl implements HnswIndex { * Constructor for loading the index from an {@link InputStream} * * @param inputStream an instance of stream to read the index bytes from - * @param resources an instance of {@link CuVSResourcesImpl} + * @param resources an instance of {@link CuVSResources} */ private HnswIndexImpl( - InputStream inputStream, CuVSResourcesImpl resources, HnswIndexParams hnswIndexParams) + InputStream inputStream, CuVSResources resources, HnswIndexParams hnswIndexParams) throws Throwable { this.hnswIndexParams = hnswIndexParams; this.resources = resources; @@ -105,8 +105,6 @@ public SearchResults search(HnswQuery query) throws Throwable { MemorySegment distancesMemorySegment = localArena.allocate(distancesSequenceLayout); MemorySegment querySeg = buildMemorySegment(localArena, queryVectors); - long cuvsRes = resources.getHandle(); - long[] queriesShape = {numQueries, vectorDimension}; MemorySegment queriesTensor = prepareTensor(localArena, querySeg, queriesShape, 2, 32, 2, 1, 1); @@ -117,21 +115,24 @@ public SearchResults search(HnswQuery query) throws Throwable { MemorySegment distancesTensor = prepareTensor(localArena, distancesMemorySegment, distancesShape, 2, 32, 2, 1, 1); - int returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); - - returnValue = - cuvsHnswSearch( - cuvsRes, - segmentFromSearchParams(localArena, query.getHnswSearchParams()), - hnswIndexReference.getMemorySegment(), - queriesTensor, - neighborsTensor, - distancesTensor); - checkCuVSError(returnValue, "cuvsHnswSearch"); - - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); + try (var resourcesAccessor = resources.access()) { + var cuvsRes = resourcesAccessor.handle(); + int returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + returnValue = + cuvsHnswSearch( + cuvsRes, + segmentFromSearchParams(localArena, query.getHnswSearchParams()), + hnswIndexReference.getMemorySegment(), + queriesTensor, + neighborsTensor, + distancesTensor); + checkCuVSError(returnValue, "cuvsHnswSearch"); + + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + } return HnswSearchResults.create( neighborsSequenceLayout, @@ -178,8 +179,6 @@ private IndexReference deserialize(InputStream inputStream) throws Throwable { inputStream.transferTo(outputStream); MemorySegment pathSeg = buildMemorySegment(localArena, tmpIndexFile.toString()); - long cuvsRes = resources.getHandle(); - var indexReference = createHnswIndex(); MemorySegment dtype = DLDataType.allocate(localArena); @@ -189,15 +188,18 @@ private IndexReference deserialize(InputStream inputStream) throws Throwable { cuvsHnswIndex.dtype(indexReference.memorySegment, dtype); - var returnValue = - cuvsHnswDeserialize( - cuvsRes, - segmentFromIndexParams(localArena, hnswIndexParams), - pathSeg, - hnswIndexParams.getVectorDimension(), - 0, - indexReference.memorySegment); - checkCuVSError(returnValue, "cuvsHnswDeserialize"); + try (var resourcesAccessor = resources.access()) { + var cuvsRes = resourcesAccessor.handle(); + var returnValue = + cuvsHnswDeserialize( + cuvsRes, + segmentFromIndexParams(localArena, hnswIndexParams), + pathSeg, + hnswIndexParams.getVectorDimension(), + 0, + indexReference.memorySegment); + checkCuVSError(returnValue, "cuvsHnswDeserialize"); + } return indexReference; @@ -227,11 +229,7 @@ private static MemorySegment segmentFromSearchParams(Arena arena, HnswSearchPara } public static HnswIndex.Builder newBuilder(CuVSResources cuvsResources) { - Objects.requireNonNull(cuvsResources); - if (!(cuvsResources instanceof CuVSResourcesImpl)) { - throw new IllegalArgumentException("Unsupported " + cuvsResources); - } - return new HnswIndexImpl.Builder((CuVSResourcesImpl) cuvsResources); + return new HnswIndexImpl.Builder(Objects.requireNonNull(cuvsResources)); } /** @@ -239,7 +237,7 @@ public static HnswIndex.Builder newBuilder(CuVSResources cuvsResources) { */ public static class Builder implements HnswIndex.Builder { - private final CuVSResourcesImpl cuvsResources; + private final CuVSResources cuvsResources; private InputStream inputStream; private HnswIndexParams hnswIndexParams; @@ -248,7 +246,7 @@ public static class Builder implements HnswIndex.Builder { * * @param cuvsResources an instance of {@link CuVSResources} */ - public Builder(CuVSResourcesImpl cuvsResources) { + public Builder(CuVSResources cuvsResources) { this.cuvsResources = cuvsResources; } diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java index 035bbe2eef..f5dcb5eff0 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/TieredIndexImpl.java @@ -28,16 +28,7 @@ import static com.nvidia.cuvs.internal.common.Util.checkCudaError; import static com.nvidia.cuvs.internal.common.Util.concatenate; import static com.nvidia.cuvs.internal.common.Util.prepareTensor; -import static com.nvidia.cuvs.internal.panama.headers_h.cudaMemcpy; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMAlloc; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsRMMFree; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsStreamSync; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsTieredIndexBuild; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsTieredIndexCreate; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsTieredIndexDestroy; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsTieredIndexExtend; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsTieredIndexSearch; -import static com.nvidia.cuvs.internal.panama.headers_h.cuvsTieredIndex_t; +import static com.nvidia.cuvs.internal.panama.headers_h.*; import com.nvidia.cuvs.CagraIndexParams; import com.nvidia.cuvs.CagraSearchParams; @@ -71,9 +62,8 @@ * @since 25.02 */ public class TieredIndexImpl implements TieredIndex { - private final float[][] vectors; private final Dataset dataset; - private final CuVSResourcesImpl resources; + private final CuVSResources resources; private final TieredIndexParams tieredIndexParameters; private final IndexReference tieredIndexReference; private boolean destroyed; @@ -82,13 +72,8 @@ public class TieredIndexImpl implements TieredIndex { * Constructor for building the index using specified dataset */ private TieredIndexImpl( - TieredIndexParams indexParameters, - float[][] vectors, - Dataset dataset, - CuVSResourcesImpl resources) - throws Throwable { + TieredIndexParams indexParameters, Dataset dataset, CuVSResources resources) { this.tieredIndexParameters = indexParameters; - this.vectors = vectors; this.dataset = dataset; this.resources = resources; this.tieredIndexReference = build(); @@ -98,22 +83,10 @@ private TieredIndexImpl( /** * Constructor for loading the index from an {@link InputStream} */ - private TieredIndexImpl(InputStream inputStream, CuVSResourcesImpl resources) throws Throwable { + private TieredIndexImpl(InputStream inputStream, CuVSResources resources) { throw new UnsupportedOperationException("Deserialization of TieredIndex is not yet supported"); } - /** - * Constructor for creating an index from an existing index reference. - */ - private TieredIndexImpl(IndexReference indexReference, CuVSResourcesImpl resources) { - this.vectors = null; - this.tieredIndexParameters = null; - this.dataset = null; - this.resources = resources; - this.tieredIndexReference = indexReference; - this.destroyed = false; - } - private void checkNotDestroyed() { if (destroyed) { throw new IllegalStateException("destroyed"); @@ -129,6 +102,10 @@ public void destroyIndex() throws Throwable { try { int returnValue = cuvsTieredIndexDestroy(tieredIndexReference.getMemorySegment()); checkCuVSError(returnValue, "cuvsTieredIndexDestroy"); + if (dataset != null) { + dataset.close(); + } + } finally { destroyed = true; } @@ -142,10 +119,10 @@ public void destroyIndex() throws Throwable { * @return an instance of {@link IndexReference} that holds the pointer to the * index */ - private IndexReference build() throws Throwable { + private IndexReference build() { try (var localArena = Arena.ofConfined()) { - long rows = dataset != null ? dataset.size() : vectors.length; - long cols = dataset != null ? dataset.dimensions() : (rows > 0 ? vectors[0].length : 0); + long rows = dataset.size(); + long cols = dataset.dimensions(); MemorySegment indexParamsMemorySegment = tieredIndexParameters != null @@ -153,48 +130,43 @@ private IndexReference build() throws Throwable { : MemorySegment.NULL; // Get host data - MemorySegment hostDataSeg = - dataset != null - ? ((DatasetImpl) dataset).asMemorySegment() - : Util.buildMemorySegment(localArena, vectors); + MemorySegment hostDataSeg = ((DatasetImpl) dataset).asMemorySegment(); - long cuvsRes = resources.getHandle(); + try (var resourceAccess = resources.access()) { + long cuvsRes = resourceAccess.handle(); - // TieredIndex REQUIRES device memory - allocate it - MemorySegment datasetD = localArena.allocate(C_POINTER); - long datasetSize = C_FLOAT_BYTE_SIZE * rows * cols; - int returnValue = cuvsRMMAlloc(cuvsRes, datasetD, datasetSize); - checkCuVSError(returnValue, "cuvsRMMAlloc"); + // TieredIndex REQUIRES device memory - allocate it + long datasetSize = C_FLOAT_BYTE_SIZE * rows * cols; + MemorySegment datasetDP = allocateRMMSegment(cuvsRes, datasetSize); - MemorySegment datasetDP = datasetD.get(C_POINTER, 0); + // Copy host to device + Util.cudaMemcpy(datasetDP, hostDataSeg, datasetSize, HOST_TO_DEVICE); - // Copy host to device - Util.cudaMemcpy(datasetDP, hostDataSeg, datasetSize, HOST_TO_DEVICE); + // Create tensor from device memory + long[] datasetShape = {rows, cols}; + MemorySegment datasetTensor = + prepareTensor(localArena, datasetDP, datasetShape, kDLFloat(), 32, 2, kDLCUDA(), 1); - // Create tensor from device memory - long[] datasetShape = {rows, cols}; - MemorySegment datasetTensor = - prepareTensor(localArena, datasetDP, datasetShape, 2, 32, 2, 2, 1); + MemorySegment index = localArena.allocate(cuvsTieredIndex_t); + var returnValue = cuvsTieredIndexCreate(index); + checkCuVSError(returnValue, "cuvsTieredIndexCreate"); - MemorySegment index = localArena.allocate(cuvsTieredIndex_t); - returnValue = cuvsTieredIndexCreate(index); - checkCuVSError(returnValue, "cuvsTieredIndexCreate"); + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); + // Extract the actual index pointer that was written by Create + MemorySegment actualIndexPtr = index.get(C_POINTER, 0); - // Extract the actual index pointer that was written by Create - MemorySegment actualIndexPtr = index.get(C_POINTER, 0); + returnValue = + cuvsTieredIndexBuild(cuvsRes, indexParamsMemorySegment, datasetTensor, actualIndexPtr); + checkCuVSError(returnValue, "cuvsTieredIndexBuild"); - returnValue = - cuvsTieredIndexBuild(cuvsRes, indexParamsMemorySegment, datasetTensor, actualIndexPtr); - checkCuVSError(returnValue, "cuvsTieredIndexBuild"); - - // Clean up device memory after build - returnValue = cuvsRMMFree(cuvsRes, datasetDP, datasetSize); - checkCuVSError(returnValue, "cuvsRMMFree"); + // Clean up device memory after build + returnValue = cuvsRMMFree(cuvsRes, datasetDP, datasetSize); + checkCuVSError(returnValue, "cuvsRMMFree"); - return new IndexReference(actualIndexPtr); + return new IndexReference(actualIndexPtr); + } } } @@ -227,126 +199,125 @@ public SearchResults search(TieredIndexQuery query) throws Throwable { // Get host query data MemorySegment hostQueriesSeg = Util.buildMemorySegment(localArena, query.getQueryVectors()); - long cuvsRes = resources.getHandle(); - - // Allocate DEVICE memory for all data - MemorySegment queriesD = localArena.allocate(C_POINTER); - MemorySegment neighborsD = localArena.allocate(C_POINTER); - MemorySegment distancesD = localArena.allocate(C_POINTER); - - long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; - long neighborsBytes = C_LONG_BYTE_SIZE * numQueries * topK; // 64-bit for tiered index - long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; - - int returnValue = cuvsRMMAlloc(cuvsRes, queriesD, queriesBytes); - checkCuVSError(returnValue, "cuvsRMMAlloc"); - returnValue = cuvsRMMAlloc(cuvsRes, neighborsD, neighborsBytes); - checkCuVSError(returnValue, "cuvsRMMAlloc"); - returnValue = cuvsRMMAlloc(cuvsRes, distancesD, distancesBytes); - checkCuVSError(returnValue, "cuvsRMMAlloc"); - - // Get device pointers - MemorySegment queriesDP = queriesD.get(C_POINTER, 0); - MemorySegment neighborsDP = neighborsD.get(C_POINTER, 0); - MemorySegment distancesDP = distancesD.get(C_POINTER, 0); - - // Copy queries from host to device - returnValue = - cudaMemcpy(queriesDP, hostQueriesSeg, queriesBytes, 1); // cudaMemcpyHostToDevice - checkCudaError(returnValue, "cudaMemcpy"); - - // Create tensors from device memory - long queriesShape[] = {numQueries, vectorDimension}; - MemorySegment queriesTensor = - prepareTensor(localArena, queriesDP, queriesShape, 2, 32, 2, 2, 1); - long neighborsShape[] = {numQueries, topK}; - MemorySegment neighborsTensor = - prepareTensor(localArena, neighborsDP, neighborsShape, 0, 64, 2, 2, 1); // 64-bit int - long distancesShape[] = {numQueries, topK}; - MemorySegment distancesTensor = - prepareTensor(localArena, distancesDP, distancesShape, 2, 32, 2, 2, 1); - - // Sync before prefilter setup - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); - - // Handle prefilter - MemorySegment prefilter = cuvsFilter.allocate(localArena); - MemorySegment prefilterD = localArena.allocate(C_POINTER); - MemorySegment prefilterDP = MemorySegment.NULL; - long prefilterBytes = 0; - - if (query.getPrefilter() != null) { - BitSet[] prefilters = new BitSet[] {query.getPrefilter()}; - BitSet concatenatedFilters = concatenate(prefilters, (int) query.getNumDocs()); - long filters[] = concatenatedFilters.toLongArray(); - MemorySegment hostPrefilterSeg = buildMemorySegment(localArena, filters); - - long prefilterDataLength = query.getNumDocs() * prefilters.length; - long prefilterShape[] = {(prefilterDataLength + 31) / 32}; - long prefilterLen = prefilterShape[0]; - prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; - - // Allocate device memory for prefilter - returnValue = cuvsRMMAlloc(cuvsRes, prefilterD, prefilterBytes); - checkCuVSError(returnValue, "cuvsRMMAlloc"); - - prefilterDP = prefilterD.get(C_POINTER, 0); - - // Copy prefilter to device - returnValue = cudaMemcpy(prefilterDP, hostPrefilterSeg, prefilterBytes, 1); + try (var resourceAccess = resources.access()) { + long cuvsRes = resourceAccess.handle(); + + long queriesBytes = C_FLOAT_BYTE_SIZE * numQueries * vectorDimension; + long neighborsBytes = C_LONG_BYTE_SIZE * numQueries * topK; // 64-bit for tiered index + long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; + + // Allocate DEVICE memory for all data + MemorySegment queriesDP = allocateRMMSegment(cuvsRes, queriesBytes); + MemorySegment neighborsDP = allocateRMMSegment(cuvsRes, neighborsBytes); + MemorySegment distancesDP = allocateRMMSegment(cuvsRes, distancesBytes); + + // Copy queries from host to device + var returnValue = + cudaMemcpy(queriesDP, hostQueriesSeg, queriesBytes, cudaMemcpyHostToDevice()); checkCudaError(returnValue, "cudaMemcpy"); - MemorySegment prefilterTensor = - prepareTensor(localArena, prefilterDP, prefilterShape, 1, 32, 1, 2, 1); + // Create tensors from device memory + long[] queriesShape = {numQueries, vectorDimension}; + MemorySegment queriesTensor = + prepareTensor(localArena, queriesDP, queriesShape, kDLFloat(), 32, 2, kDLCUDA(), 1); + long[] neighborsShape = {numQueries, topK}; + MemorySegment neighborsTensor = + prepareTensor( + localArena, + neighborsDP, + neighborsShape, + kDLInt(), + 64, + 2, + kDLCUDA(), + 1); // 64-bit int + long[] distancesShape = {numQueries, topK}; + MemorySegment distancesTensor = + prepareTensor(localArena, distancesDP, distancesShape, kDLFloat(), 32, 2, kDLCUDA(), 1); + + // Sync before prefilter setup + returnValue = cuvsStreamSync(cuvsRes); + checkCuVSError(returnValue, "cuvsStreamSync"); + + // Handle prefilter + MemorySegment prefilter = cuvsFilter.allocate(localArena); + final MemorySegment prefilterDP; + final long prefilterBytes; + + if (query.getPrefilter() != null) { + BitSet[] prefilters = new BitSet[] {query.getPrefilter()}; + BitSet concatenatedFilters = concatenate(prefilters, (int) query.getNumDocs()); + long[] filters = concatenatedFilters.toLongArray(); + MemorySegment hostPrefilterSeg = buildMemorySegment(localArena, filters); + + long prefilterDataLength = query.getNumDocs() * prefilters.length; + long[] prefilterShape = {(prefilterDataLength + 31) / 32}; + long prefilterLen = prefilterShape[0]; + prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; + + // Allocate device memory for prefilter + prefilterDP = allocateRMMSegment(cuvsRes, prefilterBytes); + + // Copy prefilter to device + checkCudaError( + cudaMemcpy(prefilterDP, hostPrefilterSeg, prefilterBytes, cudaMemcpyHostToDevice()), + "cudaMemcpy"); + + MemorySegment prefilterTensor = + prepareTensor( + localArena, prefilterDP, prefilterShape, kDLUInt(), 32, 1, kDLCUDA(), 1); + + cuvsFilter.type(prefilter, 1); // BITSET + cuvsFilter.addr(prefilter, prefilterTensor.address()); + } else { + prefilterDP = MemorySegment.NULL; + prefilterBytes = 0; + cuvsFilter.type(prefilter, 0); // NO_FILTER + cuvsFilter.addr(prefilter, 0); + } - cuvsFilter.type(prefilter, 1); // BITSET - cuvsFilter.addr(prefilter, prefilterTensor.address()); - } else { - cuvsFilter.type(prefilter, 0); // NO_FILTER - cuvsFilter.addr(prefilter, 0); - } + // Perform search + returnValue = + cuvsTieredIndexSearch( + cuvsRes, + segmentFromSearchParams(query.getCagraSearchParameters(), localArena), + tieredIndexReference.getMemorySegment(), + queriesTensor, + neighborsTensor, + distancesTensor, + prefilter); + checkCuVSError(returnValue, "cuvsTieredIndexSearch"); + + // Copy results from device to host + returnValue = + cudaMemcpy(neighborsSeg, neighborsDP, neighborsBytes, cudaMemcpyDeviceToHost()); + checkCudaError(returnValue, "cudaMemcpy"); + returnValue = + cudaMemcpy(distancesSeg, distancesDP, distancesBytes, cudaMemcpyDeviceToHost()); + checkCudaError(returnValue, "cudaMemcpy"); - // Perform search - returnValue = - cuvsTieredIndexSearch( - cuvsRes, - segmentFromSearchParams(query.getCagraSearchParameters(), localArena), - tieredIndexReference.getMemorySegment(), - queriesTensor, - neighborsTensor, - distancesTensor, - prefilter); - checkCuVSError(returnValue, "cuvsTieredIndexSearch"); - - // Copy results from device to host - returnValue = - cudaMemcpy(neighborsSeg, neighborsDP, neighborsBytes, 2); // cudaMemcpyDeviceToHost - checkCudaError(returnValue, "cudaMemcpy"); - returnValue = cudaMemcpy(distancesSeg, distancesDP, distancesBytes, 2); - checkCudaError(returnValue, "cudaMemcpy"); - - // Clean up device memory - returnValue = cuvsRMMFree(cuvsRes, queriesDP, queriesBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - returnValue = cuvsRMMFree(cuvsRes, neighborsDP, neighborsBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - returnValue = cuvsRMMFree(cuvsRes, distancesDP, distancesBytes); - checkCuVSError(returnValue, "cuvsRMMFree"); - - if (prefilterDP != MemorySegment.NULL) { - returnValue = cuvsRMMFree(cuvsRes, prefilterDP, prefilterBytes); + // Clean up device memory + returnValue = cuvsRMMFree(cuvsRes, queriesDP, queriesBytes); checkCuVSError(returnValue, "cuvsRMMFree"); - } + returnValue = cuvsRMMFree(cuvsRes, neighborsDP, neighborsBytes); + checkCuVSError(returnValue, "cuvsRMMFree"); + returnValue = cuvsRMMFree(cuvsRes, distancesDP, distancesBytes); + checkCuVSError(returnValue, "cuvsRMMFree"); + + if (prefilterDP != MemorySegment.NULL) { + returnValue = cuvsRMMFree(cuvsRes, prefilterDP, prefilterBytes); + checkCuVSError(returnValue, "cuvsRMMFree"); + } - return TieredSearchResultsImpl.create( - neighborsLayout, - distancesLayout, - neighborsSeg, - distancesSeg, - topK, - query.getMapping(), - numQueries); + return TieredSearchResultsImpl.create( + neighborsLayout, + distancesLayout, + neighborsSeg, + distancesSeg, + topK, + query.getMapping(), + numQueries); + } } } @@ -359,46 +330,39 @@ public ExtendBuilder extend() { /** * Performs the actual extend operation */ - private void performExtend(float[][] extendVectors, Dataset extendDataset) throws Throwable { + private void performExtend(Dataset extendDataset) throws Throwable { try (var localArena = Arena.ofConfined()) { - long rows = extendDataset != null ? extendDataset.size() : extendVectors.length; - long cols = extendDataset != null ? extendDataset.dimensions() : extendVectors[0].length; + long rows = extendDataset.size(); + long cols = extendDataset.dimensions(); // Get host data - MemorySegment hostDataSeg = - extendDataset != null - ? ((DatasetImpl) extendDataset).asMemorySegment() - : Util.buildMemorySegment(localArena, extendVectors); - - long cuvsRes = resources.getHandle(); + MemorySegment hostDataSeg = ((DatasetImpl) extendDataset).asMemorySegment(); - // Allocate device memory for extend data - MemorySegment datasetD = localArena.allocate(C_POINTER); - long dataSize = C_FLOAT_BYTE_SIZE * rows * cols; - int returnValue = cuvsRMMAlloc(cuvsRes, datasetD, dataSize); - checkCuVSError(returnValue, "cuvsRMMAlloc"); + try (var resourceAccess = resources.access()) { + long cuvsRes = resourceAccess.handle(); - MemorySegment datasetDP = datasetD.get(C_POINTER, 0); + // Allocate device memory for extend data + long dataSize = C_FLOAT_BYTE_SIZE * rows * cols; + MemorySegment datasetDP = allocateRMMSegment(cuvsRes, dataSize); - // Copy host to device - returnValue = cudaMemcpy(datasetDP, hostDataSeg, dataSize, 1); // cudaMemcpyHostToDevice - checkCudaError(returnValue, "cudaMemcpy"); + // Copy host to device + checkCudaError( + cudaMemcpy(datasetDP, hostDataSeg, dataSize, cudaMemcpyHostToDevice()), "cudaMemcpy"); - // Create tensor from device memory - long datasetShape[] = {rows, cols}; - MemorySegment datasetTensor = - prepareTensor(localArena, datasetDP, datasetShape, 2, 32, 2, 2, 1); + // Create tensor from device memory + long[] datasetShape = {rows, cols}; + MemorySegment datasetTensor = + prepareTensor(localArena, datasetDP, datasetShape, kDLFloat(), 32, 2, kDLCUDA(), 1); - returnValue = cuvsStreamSync(cuvsRes); - checkCuVSError(returnValue, "cuvsStreamSync"); + checkCuVSError(cuvsStreamSync(cuvsRes), "cuvsStreamSync"); - returnValue = - cuvsTieredIndexExtend(cuvsRes, datasetTensor, tieredIndexReference.getMemorySegment()); - checkCuVSError(returnValue, "cuvsTieredIndexExtend"); + checkCuVSError( + cuvsTieredIndexExtend(cuvsRes, datasetTensor, tieredIndexReference.getMemorySegment()), + "cuvsTieredIndexExtend"); - // Clean up device memory - returnValue = cuvsRMMFree(cuvsRes, datasetDP, dataSize); - checkCuVSError(returnValue, "cuvsRMMFree"); + // Clean up device memory + checkCuVSError(cuvsRMMFree(cuvsRes, datasetDP, dataSize), "cuvsRMMFree"); + } } } @@ -407,7 +371,6 @@ private void performExtend(float[][] extendVectors, Dataset extendDataset) throw */ public static class ExtendBuilder implements TieredIndex.ExtendBuilder { private final TieredIndexImpl index; - private float[][] vectors; private Dataset dataset; private ExtendBuilder(TieredIndexImpl index) { @@ -416,7 +379,7 @@ private ExtendBuilder(TieredIndexImpl index) { @Override public ExtendBuilder withDataset(float[][] vectors) { - this.vectors = vectors; + this.dataset = Dataset.ofArray(vectors); return this; } @@ -428,15 +391,11 @@ public ExtendBuilder withDataset(Dataset dataset) { @Override public void execute() throws Throwable { - if (vectors != null && dataset != null) { - throw new IllegalArgumentException( - "Please specify only one type of dataset (a float[][] or a Dataset instance)"); - } - if (vectors == null && dataset == null) { - throw new IllegalArgumentException("Must provide vectors or dataset"); + if (dataset == null) { + throw new IllegalArgumentException("Must provide a dataset"); } - index.performExtend(vectors, dataset); + index.performExtend(dataset); } } @@ -541,24 +500,20 @@ public TieredIndexType getIndexType() { */ public static TieredIndex.Builder newBuilder(CuVSResources cuvsResources) { Objects.requireNonNull(cuvsResources); - if (!(cuvsResources instanceof CuVSResourcesImpl)) { - throw new IllegalArgumentException("Unsupported " + cuvsResources); - } - return new TieredIndexImpl.Builder((CuVSResourcesImpl) cuvsResources); + return new TieredIndexImpl.Builder(cuvsResources); } /** * Builder helps configure and create an instance of {@link TieredIndex}. */ public static class Builder implements TieredIndex.Builder { - private CuVSResourcesImpl resources; - private float[][] vectors; + private final CuVSResources resources; private Dataset dataset; private TieredIndexParams params; private TieredIndexType indexType = TieredIndexType.CAGRA; private InputStream inputStream; - private Builder(CuVSResourcesImpl resources) { + private Builder(CuVSResources resources) { this.resources = resources; } @@ -570,12 +525,24 @@ public Builder from(InputStream inputStream) { @Override public Builder withDataset(float[][] vectors) { - this.vectors = vectors; + if (this.dataset != null) { + throw new IllegalArgumentException("An input dataset can only be specified once"); + } + if (vectors == null || vectors.length == 0 || vectors[0].length == 0) { + throw new IllegalArgumentException("The input vectors cannot be null or empty"); + } + this.dataset = Dataset.ofArray(vectors); return this; } @Override public Builder withDataset(Dataset dataset) { + if (this.dataset != null) { + throw new IllegalArgumentException("An input dataset can only be specified once"); + } + if (dataset == null) { + throw new IllegalArgumentException("An input dataset cannot be null"); + } this.dataset = dataset; return this; } @@ -597,17 +564,13 @@ public TieredIndex build() throws Throwable { if (inputStream != null) { return new TieredIndexImpl(inputStream, resources); } else { - if (vectors != null && dataset != null) { - throw new IllegalArgumentException( - "Please specify only one type of dataset (a float[][] or a Dataset instance)"); - } - if (vectors == null && dataset == null) { - throw new IllegalArgumentException("Must provide vectors or dataset"); + if (dataset == null) { + throw new IllegalArgumentException("Must provide a dataset"); } if (params == null) { throw new IllegalStateException("Index parameters must be provided"); } - return new TieredIndexImpl(params, vectors, dataset, resources); + return new TieredIndexImpl(params, dataset, resources); } } } @@ -618,14 +581,6 @@ public TieredIndex build() throws Throwable { public static class IndexReference { private final MemorySegment memorySegment; - /** - * Constructs TieredIndexReference and allocate the MemorySegment. - */ - protected IndexReference(CuVSResourcesImpl resources) { - // Don't allocate here - the C function will allocate - memorySegment = MemorySegment.NULL; - } - /** * Constructs TieredIndexReference with an instance of MemorySegment passed as a * parameter. diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java index 218da3d0d9..c39c3c0dfa 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceAndSearchIT.java @@ -108,7 +108,7 @@ public void testIndexingAndSearchingFlow() throws Throwable { Map.of(1, 0.15224183f, 0, 0.5906347f, 3, 0.5986643f)); for (int j = 0; j < 10; j++) { - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { BruteForceQuery cuvsQuery = new BruteForceQuery.Builder() .withTopK(3) @@ -136,7 +136,7 @@ public void testIndexingAndSearchingWithFiltering() throws Throwable { Map.of(0, 0.5906347f, 1, 0.15224195f, 3, 0.5986643f)); for (int j = 0; j < 10; j++) { - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { BruteForceQuery cuvsQuery = new BruteForceQuery.Builder() .withTopK(3) @@ -162,7 +162,7 @@ public void testIndexingAndSearchingWithFunctionMapping() throws Throwable { LongToIntFunction rotate = l -> (int) ((l + 1) % dataset.length); for (int j = 0; j < 10; j++) { - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { BruteForceQuery cuvsQuery = new BruteForceQuery.Builder() .withTopK(3) @@ -187,7 +187,7 @@ public void testIndexingAndSearchingWithListMapping() throws Throwable { LongToIntFunction rotate = SearchResults.mappingsFromList(mappings); for (int j = 0; j < 10; j++) { - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { BruteForceQuery cuvsQuery = new BruteForceQuery.Builder() .withTopK(3) diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java index 7a7c05fb0f..a998b3ee9e 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/BruteForceRandomizedIT.java @@ -106,7 +106,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro List> expected = generateExpectedResults(topK, vectors, queries, prefilters, log); // Create CuVS index and query - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { BruteForceQuery query = new BruteForceQuery.Builder() diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java index ee33de8b4e..bfbb4e2d1b 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java @@ -137,7 +137,7 @@ public void testIndexingAndSearchingFlow() throws Throwable { List> expectedResults = getExpectedResults(); int numTestsRuns = 5; - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { for (int j = 0; j < numTestsRuns; j++) { var index = indexOnce(Dataset.ofArray(dataset), resources); var indexPath = serializeOnce(index); @@ -165,7 +165,7 @@ public void testIndexingAndSearchingFlowInDifferentThreads() throws Throwable { List> expectedResults = getExpectedResults(); int numTestsRuns = 5; - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { for (int j = 0; j < numTestsRuns; j++) { runInAnotherThread( () -> { @@ -204,7 +204,7 @@ public void testIndexingAndSearchingFlowConcurrently() throws Throwable { numTestsRuns, () -> () -> { - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { var index = indexOnce(Dataset.ofArray(dataset), resources); var indexPath = serializeOnce(index); var loadedIndex = deserializeOnce(indexPath, resources); @@ -232,7 +232,7 @@ public void testIndexing() throws Throwable { numTestsRuns, () -> () -> { - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { var index = indexOnce(Dataset.ofArray(dataset), resources); index.destroyIndex(); } catch (Throwable e) { @@ -251,7 +251,7 @@ public void testSerialization() throws Throwable { numTestsRuns, () -> () -> { - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { var index = indexOnce(Dataset.ofArray(dataset), resources); var indexPath = serializeOnce(index); Files.deleteIfExists(indexPath); @@ -272,7 +272,7 @@ public void testDeserialization() throws Throwable { numTestsRuns, () -> () -> { - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { deserializeOnce(indexPath, resources).destroyIndex(); } catch (Throwable e) { throw new RuntimeException(e); @@ -283,7 +283,7 @@ public void testDeserialization() throws Throwable { } private Path createSerializedIndex(Dataset dataset) throws Throwable { - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { var index = indexOnce(dataset, resources); var indexPath = serializeOnce(index); index.destroyIndex(); @@ -303,7 +303,7 @@ public void testIndexingAndSearchingFlowWithCustomMappingFunction() throws Throw Map.of(2, 0.15224178f, 1, 0.59063464f, 0, 0.5986642f)); LongToIntFunction rotate = l -> (int) ((l + 1) % dataset.size()); - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { var index = indexOnce(dataset, resources); var indexPath = serializeOnce(index); var loadedIndex = deserializeOnce(indexPath, resources); @@ -326,7 +326,7 @@ public void testIndexingAndSearchingFlowWithCustomMappingList() throws Throwable Map.of(3, 0.15224178f, 4, 0.59063464f, 1, 0.5986642f)); LongToIntFunction rotate = SearchResults.mappingsFromList(mappings); - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { var index = indexOnce(dataset, resources); var indexPath = serializeOnce(index); var loadedIndex = deserializeOnce(indexPath, resources); @@ -362,7 +362,7 @@ public void testPrefilteringReducesResults() throws Throwable { .withMetric(CuvsDistanceType.L2Expanded) .build(); - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { CagraIndex index = CagraIndex.newBuilder(resources) .withDataset(dataset) @@ -543,7 +543,7 @@ public void testMergingIndexes() throws Throwable { Map.of(2, 0.5f, 3, 0.5f, 1, 180.5f), Map.of(0, 0.0f, 1, 2.0f, 2, 200.0f)); - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { CagraIndexParams indexParams = new CagraIndexParams.Builder() .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) @@ -633,7 +633,7 @@ public void testMergeStrategies() throws Throwable { Map.of(2, 0.5f, 3, 0.5f, 1, 180.5f), Map.of(0, 0.0f, 1, 2.0f, 2, 200.0f)); - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { CagraIndexParams indexParams = new CagraIndexParams.Builder() .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java index 3539a1868d..cb81f8464b 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java @@ -114,7 +114,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro List> expected = generateExpectedResults(topK, vectors, queries, prefilters, log); // Create CuVS index and query - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { CagraIndexParams indexParams = new CagraIndexParams.Builder() .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java new file mode 100644 index 0000000000..a58d9c04e7 --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CheckedCuVSResources.java @@ -0,0 +1,74 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import java.nio.file.Path; +import java.util.concurrent.atomic.AtomicLong; + +public class CheckedCuVSResources implements CuVSResources { + + private volatile boolean destroyed; + private final AtomicLong currentThreadId = new AtomicLong(0); + + private final CuVSResources inner; + + private CheckedCuVSResources(CuVSResources inner) { + this.inner = inner; + } + + static CuVSResources create() throws Throwable { + return new CheckedCuVSResources(CuVSResources.create()); + } + + private void checkNotDestroyed() { + if (destroyed) { + throw new IllegalStateException("Already destroyed"); + } + } + + @Override + public ScopedAccess access() { + checkNotDestroyed(); + var previousThreadId = currentThreadId.compareAndExchange(0, Thread.currentThread().threadId()); + if (previousThreadId != 0) { + throw new IllegalStateException( + "This resource is already accessed by thread [" + previousThreadId + "]"); + } + return new ScopedAccess() { + @Override + public long handle() { + checkNotDestroyed(); + return inner.access().handle(); + } + + @Override + public void close() { + currentThreadId.set(0); + } + }; + } + + @Override + public void close() { + destroyed = true; + inner.close(); + } + + @Override + public Path tempDirectory() { + return inner.tempDirectory(); + } +} diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSResourcesIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSResourcesIT.java new file mode 100644 index 0000000000..a599089dce --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CuVSResourcesIT.java @@ -0,0 +1,98 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.nvidia.cuvs; + +import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.*; + +import java.lang.invoke.MethodHandles; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import org.junit.Before; +import org.junit.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class CuVSResourcesIT extends CuVSTestCase { + + private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxAmd64()); + } + + @Test + public void testConcurrentAccessViaCheckedCuVSResourcesIsForbidden() throws Throwable { + var expectedError = + "This resource is already accessed by thread [" + Thread.currentThread().threadId() + "]"; + try (var resources = CheckedCuVSResources.create(); + var executor = Executors.newFixedThreadPool(1)) { + + try (var access1 = resources.access()) { + log.info( + "Outer access to resource {} from {}", + access1.handle(), + Thread.currentThread().threadId()); + + var exception = + assertThrows( + ExecutionException.class, + () -> { + var future = + executor.submit( + () -> { + try (var access2 = resources.access()) { + log.info( + "Nested access to resource {} from {}", + access2.handle(), + Thread.currentThread().threadId()); + log.info("Nested access finished"); + } + }); + future.get(); + }); + assertEquals(IllegalStateException.class, exception.getCause().getClass()); + assertEquals(expectedError, exception.getCause().getMessage()); + log.info("Outer access finished"); + } + } + } + + @Test + public void testSequentialAccessViaCheckedCuVSResourcesIsAllowed() throws Throwable { + try (var resources = CheckedCuVSResources.create(); + var executor = Executors.newFixedThreadPool(1)) { + + try (var access1 = resources.access()) { + log.info( + "Access 1 to resource {} from {}", access1.handle(), Thread.currentThread().threadId()); + } + + var future = + executor.submit( + () -> { + try (var access2 = resources.access()) { + log.info( + "Access 2 to resource {} from {}", + access2.handle(), + Thread.currentThread().threadId()); + } + }); + future.get(); + } + } +} diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswBuildAndSearchIT.java index e6d16e0a13..dd0926d200 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswBuildAndSearchIT.java @@ -119,7 +119,7 @@ public void testIndexingAndSearchingFlow() throws Throwable { Map.of(1, 0.15224178f, 0, 0.59063464f, 3, 0.59866416f)); for (int j = 0; j < 10; j++) { - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { HnswQuery hnswQuery = new HnswQuery.Builder() .withMapping(SearchResults.IDENTITY_MAPPING) @@ -144,7 +144,7 @@ public void testIndexingAndSearchingWithFunctionMapping() throws Throwable { LongToIntFunction rotate = l -> (int) ((l + 1) % dataset.length); for (int j = 0; j < 10; j++) { - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { HnswQuery hnswQuery = new HnswQuery.Builder() .withQueryVectors(queries) @@ -170,7 +170,7 @@ public void testIndexingAndSearchingWithListMapping() throws Throwable { LongToIntFunction rotate = SearchResults.mappingsFromList(mappings); for (int j = 0; j < 10; j++) { - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { HnswQuery hnswQuery = new HnswQuery.Builder() .withQueryVectors(queries) diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java index ded2f0019c..d2e198738b 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java @@ -98,7 +98,7 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro List> expected = generateExpectedResults(topK, vectors, queries, null, log); // Create CuVS index and query - try (CuVSResources resources = CuVSResources.create()) { + try (CuVSResources resources = CheckedCuVSResources.create()) { // Configure index parameters CagraIndexParams indexParams =