From e7a4a833ce243fe2987fb6ccd47abc99728292fa Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Thu, 1 Oct 2026 01:11:49 +0000 Subject: [PATCH 1/2] Account for accelerated HNSW host input memory --- ...ne-lucene99acceleratedhnswvectorswriter.md | 14 +- ...leratedhnswbinaryquantizedvectorswriter.md | 14 +- ...leratedhnswscalarquantizedvectorswriter.md | 14 +- ...lucene-api-com-nvidia-cuvs-lucene-utils.md | 45 +-- java/cuvs-lucene/README.md | 13 + .../nvidia/cuvs/lucene/HostInputMemory.java | 85 ++++++ .../Lucene99AcceleratedHNSWVectorsWriter.java | 54 ++-- ...ratedHNSWBinaryQuantizedVectorsWriter.java | 17 +- ...ratedHNSWScalarQuantizedVectorsWriter.java | 17 +- .../java/com/nvidia/cuvs/lucene/Utils.java | 34 --- .../TestAcceleratedHNSWHostInputMemory.java | 267 ++++++++++++++++++ .../TestAcceleratedHNSWMergeReplay.java | 75 ++++- .../cuvs/lucene/TestHostInputMemory.java | 238 ++++++++++++++++ 13 files changed, 755 insertions(+), 132 deletions(-) create mode 100644 java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/HostInputMemory.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWHostInputMemory.java create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestHostInputMemory.java diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md index 150ef20df2..7d14df7132 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md @@ -37,7 +37,7 @@ Initializes `Lucene99AcceleratedHNSWVectorsWriter` | --- | --- | | `IOException` | IOException | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:89`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:90`_ ### addField @@ -47,7 +47,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:131`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:133`_ ### flush @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:221`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:232`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:363`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:375`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:372`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:384`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:392`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:404`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:402`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:414`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:55`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md index b9ef604a4a..4c566ca2c6 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md @@ -37,7 +37,7 @@ Initializes `LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter` | --- | --- | | `IOException` | IOException | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:79`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:80`_ ### addField @@ -47,7 +47,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:127`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:129`_ ### flush @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:220`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:231`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:305`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:316`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:336`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:347`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:356`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:367`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:365`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:376`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:57`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md index fc8c021fb0..a4db8258cd 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md @@ -37,7 +37,7 @@ Initializes `LuceneAcceleratedHNSWScalarQuantizedVectorsWriter` | --- | --- | | `IOException` | IOException | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:89`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:90`_ ### addField @@ -47,7 +47,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:136`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:138`_ ### flush @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:247`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:258`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:330`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:341`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:361`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:372`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:381`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:392`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:390`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:401`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:56`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md index 05fa15a980..b992bee630 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-utils.md @@ -72,41 +72,6 @@ an instance of CuVSMatrix _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:60`_ -### createHostFloatMatrix - -```java -static CuVSHostMatrix createHostFloatMatrix(List data, int dimensions) -``` - -Builds a host-memory CuVSMatrix from a list of float vectors. - -Copies vectors directly into native host memory without creating an intermediate \{@code -float[][]\} on the heap. - -**Parameters** - -| Name | Description | -| --- | --- | -| `data` | The float vectors | -| `dimensions` | The number of float elements in each vector | - -**Returns** - -a host-memory CuVSMatrix - -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:88`_ - -### createHostByteMatrix - -```java -static CuVSHostMatrix createHostByteMatrix(List data, int bytesPerVector) -``` - -Builds a host-memory CuVSMatrix from byte vectors without first materializing the list as an -intermediate `byte[][]`. - -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:102`_ - ### createHostByteMatrixFromArray ```java @@ -115,7 +80,7 @@ static CuVSHostMatrix createHostByteMatrixFromArray(byte[][] data, int bytesPerV Builds a host-memory CuVSMatrix from a 2D byte array. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:113`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:79`_ ### nanosToMillis @@ -135,7 +100,7 @@ A utility method to convert nanoseconds to milliseconds. milliseconds -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:194`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:160`_ ### cuVSResourcesOrNull @@ -149,7 +114,7 @@ Creates an instance of CuVSResources. an instance of CuVSResources -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:203`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:169`_ ### handleThrowableWithIgnore @@ -172,7 +137,7 @@ A utility method that conditionally ignores certain throwable objects | --- | --- | | `IOException` | | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:231`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:197`_ ### info @@ -190,6 +155,6 @@ Utility to print info/debug messages via InfoStream. | `component` | the name of the index writer | | `msg` | the log message to push via the InfoStream | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:245`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:211`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java:23`_ diff --git a/java/cuvs-lucene/README.md b/java/cuvs-lucene/README.md index d49415ed07..f02bba5527 100644 --- a/java/cuvs-lucene/README.md +++ b/java/cuvs-lucene/README.md @@ -71,6 +71,19 @@ resources in a long-lived application. Class-level documentation is in the Runnable examples of CAGRA-accelerated HNSW indexing, and of indexing and searching entirely on the GPU with `CuVS2510GPUSearchCodec`, are in the [`examples/`](../../examples/java/cuvs-lucene) directory. +### Accelerated HNSW build memory + +The float, binary-quantized, and scalar-quantized HNSW writers include their primary native host input +payload in the codec writer's `ramBytesUsed()` while filling the matrix, building the index, and +writing its graph. Enable Lucene diagnostics with `IndexWriterConfig.setInfoStream(System.out)` to +see `primary_host_input_bytes` messages identifying each segment and field's allocation. + +This is not a process-memory measurement or an allocation limit. It excludes upper-layer inputs, +adjacency matrices, GPU workspace, and other temporary build storage. Lucene's public +`IndexWriter.ramBytesUsed()` uses cached indexing counters and does not automatically observe these +flush/merge allocations. On cleanup failure, the accounting scope ends without claiming that native +memory was successfully released. Flush policy, segment-size limits, and index formats are unchanged. + ## Contributing If you are interested in contributing to cuvs-lucene, please read the cuVS [Contributing guide](https://docs.nvidia.com/cuvs/developer-guide/contributing). diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/HostInputMemory.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/HostInputMemory.java new file mode 100644 index 0000000000..60573ac66b --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/HostInputMemory.java @@ -0,0 +1,85 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix; +import java.io.IOException; +import java.util.function.Supplier; +import org.apache.lucene.util.Accountable; +import org.apache.lucene.util.IOConsumer; +import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.RamUsageEstimator; + +/** + * Tracks the compact primary host input while one field is filled, built, and written. + * Each writer uses one instance, with allocation scopes run sequentially by one caller at a time. + * + *

This is payload accounting, not an allocator measurement or memory limit. It excludes + * upper-layer inputs, adjacency matrices, GPU workspace, and other temporary build storage. + * Lucene's cached IndexWriter accounting does not poll this value during flushes or merges. + */ +final class HostInputMemory implements Accountable { + private static final long SHALLOW_BYTES = + RamUsageEstimator.shallowSizeOfInstance(HostInputMemory.class); + + private final InfoStream infoStream; + private final String component; + private final String segment; + private long inputBytes; + + HostInputMemory(InfoStream infoStream, String component, String segment) { + this.infoStream = infoStream; + this.component = component; + this.segment = segment; + } + + /** The action must close the built dataset, or transfer it to an index that it closes. */ + void withMatrix( + String field, + long rows, + long columns, + CuVSMatrix.DataType type, + IOConsumer> buildAndWrite) + throws IOException { + long bytes = payloadBytes(rows, columns, type); + withAllocation(field, bytes, () -> CuVSMatrix.hostBuilder(rows, columns, type), buildAndWrite); + } + + // The factory boundary lets lifecycle tests inject allocation/cleanup failures without native + // RAM. + void withAllocation( + String field, + long bytes, + Supplier> allocate, + IOConsumer> buildAndWrite) + throws IOException { + try (CuVSMatrix.Builder builder = allocate.get()) { + // Built-in host builders allocate the complete compact matrix before returning. + inputBytes = bytes; + if (infoStream.isEnabled(component)) { + infoStream.message( + component, + "primary_host_input_bytes=" + bytes + " segment=" + segment + " field=" + field); + } + buildAndWrite.accept(builder); + } finally { + // Scope completion is not proof of deallocation if native cleanup itself failed. + inputBytes = 0; + } + } + + static long payloadBytes(long rows, long columns, CuVSMatrix.DataType type) { + if (rows < 0 || columns < 0) { + throw new IllegalArgumentException("Matrix dimensions must be nonnegative"); + } + return Math.multiplyExact(Math.multiplyExact(rows, columns), type.bytes()); + } + + @Override + public long ramBytesUsed() { + return SHALLOW_BYTES + inputBytes; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index bbafbe827a..4756f99684 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -64,6 +64,7 @@ public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; + private final HostInputMemory hostInputMemory; private IndexOutput hnswMeta = null; private IndexOutput hnswVectorIndex = null; private String vemFileName; @@ -95,6 +96,7 @@ public Lucene99AcceleratedHNSWVectorsWriter( super(); this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; + this.hostInputMemory = new HostInputMemory(infoStream, COMPONENT, state.segmentInfo.name); this.acceleratedHNSWParams = acceleratedHNSWParams; vemFileName = IndexFileNames.segmentFileName( @@ -153,8 +155,17 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) thro if (writeTrivialField(fieldInfo, vectors.size())) { return; } - CuVSMatrix dataset = Utils.createHostFloatMatrix(vectors, fieldInfo.getVectorDimension()); - writeNonTrivialField(fieldInfo, dataset); + hostInputMemory.withMatrix( + fieldInfo.name, + vectors.size(), + fieldInfo.getVectorDimension(), + CuVSMatrix.DataType.FLOAT, + builder -> { + for (float[] vector : vectors) { + builder.addVector(vector); + } + writeNonTrivialField(fieldInfo, builder.build()); + }); } /** @@ -302,34 +313,35 @@ private void vectorBasedMerge(FieldInfo fieldInfo, MergeState mergeState) throws FloatVectorValues mergedVectors = KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState); int dims = fieldInfo.getVectorDimension(); - CuVSHostMatrix dataset = - buildMergedDataset( - mergedVectors, size, CuVSMatrix.hostBuilder(size, dims, CuVSMatrix.DataType.FLOAT)); - writeNonTrivialField(fieldInfo, dataset); + hostInputMemory.withMatrix( + fieldInfo.name, + size, + dims, + CuVSMatrix.DataType.FLOAT, + builder -> + writeNonTrivialField(fieldInfo, buildMergedDataset(mergedVectors, size, builder))); } catch (Throwable t) { Utils.handleThrowable(t); } } - /* Replays merged vectors into a builder that remains responsible for storage until build. */ + /* Borrows a builder from the caller, which owns cleanup on both success and failure. */ static CuVSHostMatrix buildMergedDataset( FloatVectorValues mergedVectors, int expectedSize, CuVSMatrix.Builder builder) throws IOException { - try (builder) { - KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); - int replayed = 0; - for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { - if (replayed == expectedSize) { - throw mergeReplayMismatch(expectedSize, (long) replayed + 1L, true); - } - builder.addVector(mergedVectors.vectorValue(it.index())); - replayed = Math.incrementExact(replayed); - } - if (replayed != expectedSize) { - throw mergeReplayMismatch(expectedSize, replayed, false); + KnnVectorValues.DocIndexIterator it = mergedVectors.iterator(); + int replayed = 0; + for (int doc = it.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = it.nextDoc()) { + if (replayed == expectedSize) { + throw mergeReplayMismatch(expectedSize, (long) replayed + 1L, true); } - return builder.build(); + builder.addVector(mergedVectors.vectorValue(it.index())); + replayed = Math.incrementExact(replayed); + } + if (replayed != expectedSize) { + throw mergeReplayMismatch(expectedSize, replayed, false); } + return builder.build(); } private static IOException mergeReplayMismatch(int expected, long observed, boolean lowerBound) { @@ -402,7 +414,7 @@ public void close() throws IOException { */ @Override public long ramBytesUsed() { - long total = SHALLOW_RAM_BYTES_USED; + long total = SHALLOW_RAM_BYTES_USED + hostInputMemory.ramBytesUsed(); for (var field : fields) { total += field.ramBytesUsed(); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 9efcbec3e2..c7399a10a9 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -63,6 +63,7 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter extends KnnVector private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; + private final HostInputMemory hostInputMemory; private final AcceleratedHNSWParams acceleratedHNSWParams; private IndexOutput hnswMeta = null, hnswVectorIndex = null; private boolean finished; @@ -86,6 +87,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter( this.acceleratedHNSWParams = acceleratedHNSWParams; this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; + this.hostInputMemory = new HostInputMemory(infoStream, COMPONENT, state.segmentInfo.name); vemFileName = IndexFileNames.segmentFileName( @@ -154,8 +156,17 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throw try { int dimensions = fieldInfo.getVectorDimension(); int bytesPerVector = (dimensions + 7) / 8; - CuVSMatrix dataset = Utils.createHostByteMatrix(vectors, bytesPerVector); - writeNonTrivialField(fieldInfo, dataset); + hostInputMemory.withMatrix( + fieldInfo.name, + size, + bytesPerVector, + CuVSMatrix.DataType.BYTE, + builder -> { + for (byte[] vector : vectors) { + builder.addVector(vector); + } + writeNonTrivialField(fieldInfo, builder.build()); + }); } catch (Throwable t) { throw Utils.handleThrowable(t); } @@ -365,7 +376,7 @@ public void close() throws IOException { */ @Override public long ramBytesUsed() { - long total = SHALLOW_RAM_BYTES_USED; + long total = SHALLOW_RAM_BYTES_USED + hostInputMemory.ramBytesUsed(); for (var field : fields) { total += field.ramBytesUsed(); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 60cc338e4d..bb131fdcf5 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -64,6 +64,7 @@ public class LuceneAcceleratedHNSWScalarQuantizedVectorsWriter extends KnnVector private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; + private final HostInputMemory hostInputMemory; private final AcceleratedHNSWParams acceleratedHNSWParams; private IndexOutput hnswMeta = null, hnswVectorIndex = null; private boolean finished; @@ -96,6 +97,7 @@ public LuceneAcceleratedHNSWScalarQuantizedVectorsWriter( this.acceleratedHNSWParams = acceleratedHNSWParams; this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; + this.hostInputMemory = new HostInputMemory(infoStream, COMPONENT, state.segmentInfo.name); vemFileName = IndexFileNames.segmentFileName( @@ -181,8 +183,17 @@ private void writeFieldInternal(FieldInfo fieldInfo, List vectors) throws IOE } // Create CuVSMatrix with BYTE data type (unsigned bytes) - CuVSMatrix dataset = Utils.createHostByteMatrix(unsignedVectors, dimensions); - writeNonTrivialField(fieldInfo, dataset); + hostInputMemory.withMatrix( + fieldInfo.name, + size, + dimensions, + CuVSMatrix.DataType.BYTE, + builder -> { + for (byte[] vector : unsignedVectors) { + builder.addVector(vector); + } + writeNonTrivialField(fieldInfo, builder.build()); + }); } catch (Throwable t) { throw Utils.handleThrowable(t); } @@ -390,7 +401,7 @@ public void close() throws IOException { */ @Override public long ramBytesUsed() { - long total = SHALLOW_RAM_BYTES_USED; + long total = SHALLOW_RAM_BYTES_USED + hostInputMemory.ramBytesUsed(); for (var field : fields) { total += field.ramBytesUsed(); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java index 4d348f6d75..450b3695f1 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Utils.java @@ -76,40 +76,6 @@ static CuVSMatrix createFloatMatrix(List data, int dimensions, CuVSReso return builder.build(); } - /** - * Builds a host-memory CuVSMatrix from a list of float vectors. - * - *

Copies vectors directly into native host memory without creating an intermediate {@code - * float[][]} on the heap. - * - * @param data The float vectors - * @param dimensions The number of float elements in each vector - * @return a host-memory CuVSMatrix - */ - static CuVSHostMatrix createHostFloatMatrix(List data, int dimensions) { - try (CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.size(), dimensions, CuVSMatrix.DataType.FLOAT)) { - for (float[] vector : data) { - builder.addVector(vector); - } - return builder.build(); - } - } - - /** - * Builds a host-memory CuVSMatrix from byte vectors without first materializing the list as an - * intermediate {@code byte[][]}. - */ - static CuVSHostMatrix createHostByteMatrix(List data, int bytesPerVector) { - try (CuVSMatrix.Builder builder = - CuVSMatrix.hostBuilder(data.size(), bytesPerVector, CuVSMatrix.DataType.BYTE)) { - for (byte[] vector : data) { - builder.addVector(vector); - } - return builder.build(); - } - } - /** Builds a host-memory CuVSMatrix from a 2D byte array. */ static CuVSHostMatrix createHostByteMatrixFromArray(byte[][] data, int bytesPerVector) { try (CuVSMatrix.Builder builder = diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWHostInputMemory.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWHostInputMemory.java new file mode 100644 index 0000000000..44952f4f73 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWHostInputMemory.java @@ -0,0 +1,267 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.carrotsearch.randomizedtesting.annotations.Name; +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Set; +import org.apache.lucene.codecs.KnnFieldVectorsWriter; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.KnnVectorsWriter; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.MergeState; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.SegmentReadState; +import org.apache.lucene.index.SegmentWriteState; +import org.apache.lucene.index.SerialMergeScheduler; +import org.apache.lucene.index.Sorter; +import org.apache.lucene.index.TieredMergePolicy; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.IOConsumer; +import org.apache.lucene.util.InfoStream; + +/** Observes the real codec writer, not IndexWriter's cached buffering counters. Requires cuVS. */ +public class TestAcceleratedHNSWHostInputMemory extends LuceneTestCase { + private static final int VECTORS_PER_SEGMENT = 256; + private static final String[] FIELDS = {"embedding", "secondary_embedding"}; + private final KnnVectorsFormat format; + private final int dimensions; + private final int payloadBytesPerVector; + + public TestAcceleratedHNSWHostInputMemory( + @Name("format") KnnVectorsFormat format, + @Name("dimensions") int dimensions, + @Name("payloadBytesPerVector") int payloadBytesPerVector) { + this.format = format; + this.dimensions = dimensions; + this.payloadBytesPerVector = payloadBytesPerVector; + } + + @ParametersFactory + public static List parameters() { + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder() + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withGraphDegree(32) + .withIntermediateGraphDegree(64) + .withHNSWLayer(1) + .build(); + return Arrays.asList( + new Object[][] { + {new Lucene99AcceleratedHNSWVectorsFormat(params), 128, 512}, + {new LuceneAcceleratedHNSWScalarQuantizedVectorsFormat(params), 128, 128}, + // 129 binary dimensions deliberately exercise the partial final byte. + {new LuceneAcceleratedHNSWBinaryQuantizedVectorsFormat(params), 129, 17} + }); + } + + public void testFlushAndMergeCountEachFieldsNativeInput() throws Exception { + AccountingObserver observer = new AccountingObserver(); + try (Directory directory = newDirectory(); + IndexWriter writer = new IndexWriter(directory, config(observer))) { + addDocuments(writer, 0, VECTORS_PER_SEGMENT); + writer.commit(); + addDocuments(writer, VECTORS_PER_SEGMENT, VECTORS_PER_SEGMENT); + writer.commit(); + try (DirectoryReader reader = DirectoryReader.open(writer)) { + assertEquals(2, reader.leaves().size()); + } + writer.getConfig().setMergePolicy(new TieredMergePolicy()); + writer.forceMerge(1); + writer.commit(); + + long segmentBytes = (long) VECTORS_PER_SEGMENT * payloadBytesPerVector; + List bytesPerSegment = List.of(segmentBytes, segmentBytes, 2 * segmentBytes); + assertEquals(bytesPerSegment.size(), observer.segments.size()); + List expected = new ArrayList<>(); + for (int segment = 0; segment < bytesPerSegment.size(); segment++) { + for (String field : FIELDS) { + expected.add( + new InputAllocation( + observer.segments.get(segment), field, bytesPerSegment.get(segment))); + } + } + // Field traversal order is not a contract; identity, size, and exactly-once reporting are. + assertEquals(expected.size(), observer.allocations.size()); + assertEquals(Set.copyOf(expected), Set.copyOf(observer.allocations)); + try (DirectoryReader reader = DirectoryReader.open(writer)) { + assertEquals(1, reader.leaves().size()); + assertEquals(2 * VECTORS_PER_SEGMENT, reader.numDocs()); + for (String field : FIELDS) { + assertEquals( + reader.numDocs(), getOnlyLeafReader(reader).getFloatVectorValues(field).size()); + } + } + TestUtil.checkIndex(directory); + } + } + + public void testSingletonFieldsDoNotAllocatePrimaryHostInput() throws Exception { + AccountingObserver observer = new AccountingObserver(); + try (Directory directory = newDirectory(); + IndexWriter writer = new IndexWriter(directory, config(observer))) { + addDocuments(writer, 0, 1); + writer.commit(); + assertTrue(observer.allocations.isEmpty()); + assertTrue("The accelerated writer must actually flush", observer.completedWrites > 0); + } + } + + private IndexWriterConfig config(AccountingObserver observer) { + assertTrue( + "These accounting integration tests require GPU/cuVS support", + ThreadLocalCuVSResourcesProvider.isSupported()); + return newIndexWriterConfig() + .setCodec(TestUtil.alwaysKnnVectorsFormat(observingFormat(observer))) + .setInfoStream(observer) + .setMergePolicy(NoMergePolicy.INSTANCE) + .setMergeScheduler(new SerialMergeScheduler()) + .setMaxBufferedDocs(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setRAMBufferSizeMB(256); + } + + private void addDocuments(IndexWriter writer, int firstId, int count) throws IOException { + for (int id = firstId; id < firstId + count; id++) { + Document document = new Document(); + for (String field : FIELDS) { + float[] vector = new float[dimensions]; + for (int dimension = 0; dimension < dimensions; dimension++) { + vector[dimension] = (float) Math.sin((id + 1.0) * (dimension + 1.0)); + } + document.add(new KnnFloatVectorField(field, vector)); + } + writer.addDocument(document); + } + } + + private KnnVectorsFormat observingFormat(AccountingObserver observer) { + return new KnnVectorsFormat(format.getName()) { + @Override + public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException { + KnnVectorsWriter delegate = format.fieldsWriter(state); + assertTrue( + "CPU fallback must not satisfy accounting tests", + delegate instanceof Lucene99AcceleratedHNSWVectorsWriter + || delegate instanceof LuceneAcceleratedHNSWScalarQuantizedVectorsWriter + || delegate instanceof LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter); + // Capture identities independently of the diagnostic text being checked. + observer.segments.add(state.segmentInfo.name); + return new ObservedWriter(delegate, observer, state.segmentInfo.name); + } + + @Override + public KnnVectorsReader fieldsReader(SegmentReadState state) throws IOException { + return format.fieldsReader(state); + } + + @Override + public int getMaxDimensions(String fieldName) { + return format.getMaxDimensions(fieldName); + } + }; + } + + private record InputAllocation(String segment, String field, long bytes) {} + + private static final class AccountingObserver extends InfoStream { + private final List allocations = new ArrayList<>(); + private final List segments = new ArrayList<>(); + private ObservedWriter active; + private long baseline; + private int completedWrites; + + @Override + public void message(String component, String message) { + if (!message.startsWith("primary_host_input_bytes=")) { + return; + } + assertNotNull("Allocation must occur inside flush or merge", active); + long bytes = + Long.parseLong(message.substring(message.indexOf('=') + 1, message.indexOf(' '))); + assertTrue(message.contains(" segment=" + active.segment + " field=")); + String field = message.substring(message.indexOf(" field=") + " field=".length()); + assertEquals( + "Codec accounting must include this allocation exactly once", + baseline + bytes, + active.delegate.ramBytesUsed()); + allocations.add(new InputAllocation(active.segment, field, bytes)); + } + + @Override + public boolean isEnabled(String component) { + return true; + } + + @Override + public void close() {} + } + + /** Captures a baseline immediately before each synchronous codec operation. */ + private static final class ObservedWriter extends KnnVectorsWriter { + private final KnnVectorsWriter delegate; + private final AccountingObserver observer; + private final String segment; + + ObservedWriter(KnnVectorsWriter delegate, AccountingObserver observer, String segment) { + this.delegate = delegate; + this.observer = observer; + this.segment = segment; + } + + private void observe(IOConsumer operation) throws IOException { + observer.active = this; + observer.baseline = delegate.ramBytesUsed(); + operation.accept(delegate); + assertEquals( + "Completed scope must not retain native input bytes", + observer.baseline, + delegate.ramBytesUsed()); + observer.completedWrites++; + observer.active = null; + } + + @Override + public KnnFieldVectorsWriter addField(FieldInfo info) throws IOException { + return delegate.addField(info); + } + + @Override + public void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException { + observe(writer -> writer.flush(maxDoc, sortMap)); + } + + @Override + public void mergeOneField(FieldInfo info, MergeState state) throws IOException { + observe(writer -> writer.mergeOneField(info, state)); + } + + @Override + public void finish() throws IOException { + delegate.finish(); + } + + @Override + public void close() throws IOException { + delegate.close(); + } + + @Override + public long ramBytesUsed() { + return delegate.ramBytesUsed(); + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMergeReplay.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMergeReplay.java index c2bd07b87f..8a6597daf4 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMergeReplay.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMergeReplay.java @@ -7,10 +7,12 @@ import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; import java.io.IOException; +import java.lang.reflect.Proxy; import java.util.List; import org.apache.lucene.index.FloatVectorValues; import org.apache.lucene.index.KnnVectorValues; import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.InfoStream; public class TestAcceleratedHNSWMergeReplay extends LuceneTestCase { @@ -19,9 +21,7 @@ public void testUnderflowClosesBuilderBeforeBuild() { FloatVectorValues values = FloatVectorValues.fromFloats(List.of(new float[] {1f}), 1); IOException failure = - expectThrows( - IOException.class, - () -> Lucene99AcceleratedHNSWVectorsWriter.buildMergedDataset(values, 2, builder)); + expectThrows(IOException.class, () -> replayWithinAccountingScope(values, 2, builder)); assertEquals( "Merged vector count changed between passes: expected 2, observed 1", failure.getMessage()); @@ -36,9 +36,7 @@ public void testOverflowClosesBuilderBeforeBuild() { FloatVectorValues.fromFloats(List.of(new float[] {1f}, new float[] {2f}), 1); IOException failure = - expectThrows( - IOException.class, - () -> Lucene99AcceleratedHNSWVectorsWriter.buildMergedDataset(values, 1, builder)); + expectThrows(IOException.class, () -> replayWithinAccountingScope(values, 1, builder)); assertEquals( "Merged vector count changed between passes: expected 1, observed at least 2", @@ -55,9 +53,7 @@ public void testReplayFailureRemainsPrimaryWhenBuilderCloseFails() { FloatVectorValues values = failingValues(replayFailure); IOException thrown = - expectThrows( - IOException.class, - () -> Lucene99AcceleratedHNSWVectorsWriter.buildMergedDataset(values, 1, builder)); + expectThrows(IOException.class, () -> replayWithinAccountingScope(values, 1, builder)); assertSame(replayFailure, thrown); assertArrayEquals(new Throwable[] {closeFailure}, thrown.getSuppressed()); @@ -66,6 +62,58 @@ public void testReplayFailureRemainsPrimaryWhenBuilderCloseFails() { assertEquals(1, builder.closeCalls); } + public void testSuccessfulReplayClosesBuilderExactlyOnce() throws IOException { + TrackingBuilder builder = new TrackingBuilder(); + int[] datasetCloses = {0}; + // Replay only passes this returned handle to its owner. No native allocation is needed. + builder.builtMatrix = + (CuVSHostMatrix) + Proxy.newProxyInstance( + CuVSHostMatrix.class.getClassLoader(), + new Class[] {CuVSHostMatrix.class}, + (proxy, method, arguments) -> { + if (method.getName().equals("close")) { + datasetCloses[0]++; + return null; + } + throw new AssertionError( + "Replay must not inspect the returned matrix: " + method.getName()); + }); + FloatVectorValues values = FloatVectorValues.fromFloats(List.of(new float[] {1f}), 1); + + replayWithinAccountingScope(values, 1, builder); + + assertEquals(1, builder.addCalls); + assertEquals(1, builder.buildCalls); + assertEquals(1, builder.closeCalls); + assertEquals(1, datasetCloses[0]); + } + + private static void replayWithinAccountingScope( + FloatVectorValues values, int expectedSize, TrackingBuilder builder) throws IOException { + HostInputMemory memory = new HostInputMemory(InfoStream.NO_OUTPUT, "test", "_0"); + long baseline = memory.ramBytesUsed(); + long payloadBytes = (long) expectedSize * values.dimension() * Float.BYTES; + try { + memory.withAllocation( + "embedding", + payloadBytes, + () -> builder, + allocated -> { + try (CuVSHostMatrix dataset = + Lucene99AcceleratedHNSWVectorsWriter.buildMergedDataset( + values, expectedSize, allocated)) { + assertSame(builder.builtMatrix, dataset); + assertEquals(baseline + payloadBytes, memory.ramBytesUsed()); + } catch (Exception e) { + throw Utils.handleThrowable(e); + } + }); + } finally { + assertEquals(baseline, memory.ramBytesUsed()); + } + } + private static FloatVectorValues failingValues(IOException failure) { FloatVectorValues delegate = FloatVectorValues.fromFloats(List.of(new float[] {1f}), 1); return new FloatVectorValues() { @@ -101,6 +149,7 @@ private static final class TrackingBuilder implements CuVSMatrix.Builder HostInputMemory.payloadBytes(Long.MAX_VALUE, 2, CuVSMatrix.DataType.BYTE)); + expectThrows( + ArithmeticException.class, + () -> HostInputMemory.payloadBytes(Long.MAX_VALUE / 2, 1, CuVSMatrix.DataType.FLOAT)); + expectThrows( + IllegalArgumentException.class, + () -> HostInputMemory.payloadBytes(-1, 128, CuVSMatrix.DataType.FLOAT)); + } + + public void testPayloadAboveTwoGiBIsCountedUntilCleanupWithoutAllocatingIt() throws Exception { + HostInputMemory memory = newMemory(); + long baseline = memory.ramBytesUsed(); + long threeGiB = HostInputMemory.payloadBytes(3L * 1024 * 1024, 256, CuVSMatrix.DataType.FLOAT); + assertEquals(3L * 1024 * 1024 * 1024, threeGiB); + + // Repeat with a different field size to catch stale or cumulative accounting. + for (long bytes : new long[] {threeGiB, 512}) { + TrackingBuilder builder = + new TrackingBuilder(() -> assertEquals(baseline + bytes, memory.ramBytesUsed())); + memory.withAllocation( + "embedding", + bytes, + () -> { + assertEquals(baseline, memory.ramBytesUsed()); + return builder; + }, + allocated -> { + assertSame(builder, allocated); + assertEquals(baseline + bytes, memory.ramBytesUsed()); + allocated.addVector(new float[] {1f}); + }); + assertEquals(1, builder.addCalls); + assertEquals(1, builder.closeCalls); + assertEquals(baseline, memory.ramBytesUsed()); + } + } + + public void testFailedAllocationNeverPublishesPayload() { + HostInputMemory memory = newMemory(); + long baseline = memory.ramBytesUsed(); + OutOfMemoryError allocationFailure = + new OutOfMemoryError("simulated native allocation failure"); + + OutOfMemoryError thrown = + expectThrows( + OutOfMemoryError.class, + () -> + memory.withAllocation( + "embedding", + 512, + () -> { + throw allocationFailure; + }, + builder -> fail("No builder exists after failed allocation"))); + + assertSame(allocationFailure, thrown); + assertEquals(baseline, memory.ramBytesUsed()); + } + + public void testBodyFailureRemainsPrimaryWhenBuilderCleanupFails() { + HostInputMemory memory = newMemory(); + long baseline = memory.ramBytesUsed(); + IOException bodyFailure = new IOException("vector replay failed"); + IllegalStateException closeFailure = new IllegalStateException("builder cleanup failed"); + TrackingBuilder builder = + new TrackingBuilder( + () -> { + throw closeFailure; + }); + + IOException thrown = + expectThrows( + IOException.class, + () -> + memory.withAllocation( + "embedding", + 512, + () -> builder, + allocated -> { + throw bodyFailure; + })); + + assertSame(bodyFailure, thrown); + assertArrayEquals(new Throwable[] {closeFailure}, thrown.getSuppressed()); + assertEquals(1, builder.closeCalls); + assertEquals(baseline, memory.ramBytesUsed()); + } + + public void testBuildFailureClosesBuilderAndClearsAccounting() { + HostInputMemory memory = newMemory(); + long baseline = memory.ramBytesUsed(); + TrackingBuilder builder = new TrackingBuilder(() -> {}); + + IllegalStateException thrown = + expectThrows( + IllegalStateException.class, + () -> + memory.withAllocation( + "embedding", 512, () -> builder, allocated -> allocated.build())); + + assertSame(builder.buildFailure, thrown); + assertEquals(1, builder.closeCalls); + assertEquals(baseline, memory.ramBytesUsed()); + } + + public void testOwnershipTransferDoesNotDoubleCountPayload() throws Exception { + HostInputMemory memory = newMemory(); + long baseline = memory.ramBytesUsed(); + int[] indexCloses = {0}; + memory.withAllocation( + "embedding", + 512, + () -> new TrackingBuilder(() -> {}), + builder -> { + try (Utils.OwnedIndex owned = + Utils.ownDataset(() -> fail("Index owns the transferred dataset"))) { + owned.transferTo( + () -> { + assertEquals(baseline + 512, memory.ramBytesUsed()); + indexCloses[0]++; + }); + assertEquals(baseline + 512, memory.ramBytesUsed()); + } catch (Exception e) { + throw new IOException(e); + } + }); + assertEquals(1, indexCloses[0]); + assertEquals(baseline, memory.ramBytesUsed()); + } + + public void testDiagnosticFailureStillClosesAllocatedBuilder() { + IllegalStateException diagnosticFailure = new IllegalStateException("diagnostic sink failed"); + InfoStream failingLog = + new InfoStream() { + @Override + public void message(String component, String message) { + throw diagnosticFailure; + } + + @Override + public boolean isEnabled(String component) { + return true; + } + + @Override + public void close() {} + }; + HostInputMemory memory = new HostInputMemory(failingLog, "test", "_0"); + long baseline = memory.ramBytesUsed(); + TrackingBuilder builder = new TrackingBuilder(() -> {}); + + IllegalStateException thrown = + expectThrows( + IllegalStateException.class, + () -> + memory.withAllocation( + "embedding", + 512, + () -> builder, + allocated -> fail("Diagnostic failure prevents population"))); + + assertSame(diagnosticFailure, thrown); + assertEquals(1, builder.closeCalls); + assertEquals(baseline, memory.ramBytesUsed()); + } + + private static HostInputMemory newMemory() { + return new HostInputMemory(InfoStream.NO_OUTPUT, "test", "_0"); + } + + /** Only supplies the builder lifecycle; no native allocation or successful matrix build occurs. */ + private static final class TrackingBuilder implements CuVSMatrix.Builder { + private final Runnable onClose; + private final IllegalStateException buildFailure = + new IllegalStateException("matrix build failed"); + private int addCalls; + private int closeCalls; + + TrackingBuilder(Runnable onClose) { + this.onClose = onClose; + } + + @Override + public void addVector(float[] vector) { + addCalls++; + } + + @Override + public void addVector(byte[] vector) { + throw new AssertionError("unexpected byte vector"); + } + + @Override + public void addVector(int[] vector) { + throw new AssertionError("unexpected int vector"); + } + + @Override + public void addVector(short[] vector) { + throw new AssertionError("unexpected short vector"); + } + + @Override + public CuVSHostMatrix build() { + throw buildFailure; + } + + @Override + public void close() { + closeCalls++; + onClose.run(); + } + } +} From f63c9892d2be300708ddc42f0bf68f6033e0ab12 Mon Sep 17 00:00:00 2001 From: nvzm123 Date: Thu, 1 Oct 2026 02:15:51 +0000 Subject: [PATCH 2/2] Simplify host input memory accounting tests --- .../cuvs/lucene/TestHostInputMemory.java | 56 ++----------------- 1 file changed, 5 insertions(+), 51 deletions(-) diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestHostInputMemory.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestHostInputMemory.java index f6c31462e9..b359fae3fb 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestHostInputMemory.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestHostInputMemory.java @@ -6,7 +6,6 @@ import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; -import java.io.IOException; import org.apache.lucene.tests.util.LuceneTestCase; import org.apache.lucene.util.InfoStream; @@ -29,6 +28,9 @@ public void testPayloadArithmeticRejectsOverflowAndNegativeDimensions() { expectThrows( IllegalArgumentException.class, () -> HostInputMemory.payloadBytes(-1, 128, CuVSMatrix.DataType.FLOAT)); + expectThrows( + IllegalArgumentException.class, + () -> HostInputMemory.payloadBytes(1, -1, CuVSMatrix.DataType.FLOAT)); } public void testPayloadAboveTwoGiBIsCountedUntilCleanupWithoutAllocatingIt() throws Exception { @@ -81,10 +83,9 @@ public void testFailedAllocationNeverPublishesPayload() { assertEquals(baseline, memory.ramBytesUsed()); } - public void testBodyFailureRemainsPrimaryWhenBuilderCleanupFails() { + public void testBuildFailureRemainsPrimaryWhenBuilderCleanupFails() { HostInputMemory memory = newMemory(); long baseline = memory.ramBytesUsed(); - IOException bodyFailure = new IOException("vector replay failed"); IllegalStateException closeFailure = new IllegalStateException("builder cleanup failed"); TrackingBuilder builder = new TrackingBuilder( @@ -92,29 +93,6 @@ public void testBodyFailureRemainsPrimaryWhenBuilderCleanupFails() { throw closeFailure; }); - IOException thrown = - expectThrows( - IOException.class, - () -> - memory.withAllocation( - "embedding", - 512, - () -> builder, - allocated -> { - throw bodyFailure; - })); - - assertSame(bodyFailure, thrown); - assertArrayEquals(new Throwable[] {closeFailure}, thrown.getSuppressed()); - assertEquals(1, builder.closeCalls); - assertEquals(baseline, memory.ramBytesUsed()); - } - - public void testBuildFailureClosesBuilderAndClearsAccounting() { - HostInputMemory memory = newMemory(); - long baseline = memory.ramBytesUsed(); - TrackingBuilder builder = new TrackingBuilder(() -> {}); - IllegalStateException thrown = expectThrows( IllegalStateException.class, @@ -123,35 +101,11 @@ public void testBuildFailureClosesBuilderAndClearsAccounting() { "embedding", 512, () -> builder, allocated -> allocated.build())); assertSame(builder.buildFailure, thrown); + assertArrayEquals(new Throwable[] {closeFailure}, thrown.getSuppressed()); assertEquals(1, builder.closeCalls); assertEquals(baseline, memory.ramBytesUsed()); } - public void testOwnershipTransferDoesNotDoubleCountPayload() throws Exception { - HostInputMemory memory = newMemory(); - long baseline = memory.ramBytesUsed(); - int[] indexCloses = {0}; - memory.withAllocation( - "embedding", - 512, - () -> new TrackingBuilder(() -> {}), - builder -> { - try (Utils.OwnedIndex owned = - Utils.ownDataset(() -> fail("Index owns the transferred dataset"))) { - owned.transferTo( - () -> { - assertEquals(baseline + 512, memory.ramBytesUsed()); - indexCloses[0]++; - }); - assertEquals(baseline + 512, memory.ramBytesUsed()); - } catch (Exception e) { - throw new IOException(e); - } - }); - assertEquals(1, indexCloses[0]); - assertEquals(baseline, memory.ramBytesUsed()); - } - public void testDiagnosticFailureStillClosesAllocatedBuilder() { IllegalStateException diagnosticFailure = new IllegalStateException("diagnostic sink failed"); InfoStream failingLog =