diff --git a/ci/test_lucene.sh b/ci/test_lucene.sh index 5478a834ae..a07488f9b1 100755 --- a/ci/test_lucene.sh +++ b/ci/test_lucene.sh @@ -11,6 +11,9 @@ set +e rapids-logger "Check GPU usage" nvidia-smi +# This GPU lane must execute the real cuVS writer-path sentinel rather than skip it. +export CUVS_TESTS_REQUIRE_GPU=1 + rapids-logger "Run cuvs-lucene build and tests" RAPIDS_CUDA_MAJOR="${RAPIDS_CUDA_VERSION%%.*}" diff --git a/ci/test_lucene_prebuilt.sh b/ci/test_lucene_prebuilt.sh index 477ff75ad3..05f384b910 100755 --- a/ci/test_lucene_prebuilt.sh +++ b/ci/test_lucene_prebuilt.sh @@ -33,6 +33,9 @@ fi rapids-logger "Check GPU usage" nvidia-smi +# This GPU lane must execute the real cuVS writer-path sentinel rather than skip it. +export CUVS_TESTS_REQUIRE_GPU=1 + rapids-logger "Configuring conda strict channel priority" conda config --set channel_priority strict diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md index d3147c842e..f4ca78f79f 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md @@ -18,14 +18,28 @@ public class AcceleratedHNSWParams public int getWriterThreads() ``` -Get the maximum thread count for cuVS writes and accelerated-HNSW graph materialization and -serialization. +Get the native cuVS writer threads parameter. **Returns** -maximum writer and graph-processing thread count +cuVS writer threads parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:148`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:154`_ + +### getGraphThreads + +```java +public int getGraphThreads() +``` + +Get the maximum threads per HNSW graph materialization or serialization operation. The count +includes the calling thread; shared helper capacity may reduce actual concurrency. + +**Returns** + +HNSW graph processing threads parameter + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:164`_ ### getIntermediateGraphDegree @@ -39,7 +53,7 @@ Get the intermediate graph degree the graph degree parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:157`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:173`_ ### getGraphdegree @@ -53,7 +67,7 @@ Get the graph degree the graph degree parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:166`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:182`_ ### getHnswLayers @@ -67,7 +81,7 @@ Get the number of HNSW layers the number of HNSW layers -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:175`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:191`_ ### getMaxConn @@ -81,7 +95,7 @@ Get the max connection parameter the max connection parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:184`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:200`_ ### getBeamWidth @@ -95,7 +109,7 @@ Get the beam width parameter the beam width parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:193`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:209`_ ### getCagraGraphBuildAlgo @@ -109,7 +123,7 @@ Get the CAGRA graph build algorithm the CAGRA graph build algorithm -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:202`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:218`_ ### getCuVSIvfPqParams @@ -123,7 +137,7 @@ Get the instance of `CuVSIvfPqParams` the instance of `CuVSIvfPqParams` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:211`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:227`_ ### getNumMergeWorkers @@ -137,7 +151,7 @@ Get the number of merge workers set to be used in the fallback mechanism the number of merge workers -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:220`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:236`_ ### getMergeExec @@ -151,7 +165,7 @@ Get the instance of the `ExecutorService` to be used in the fallback mechanism the instance of the `ExecutorService` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:229`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:245`_ ### getStrategy @@ -168,7 +182,7 @@ When CUSTOM is chosen, the build algorithm and its parameters (either defaults o get the chosen `Strategy` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:241`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:257`_ ### getCuvsDistanceType @@ -182,7 +196,7 @@ Get the cuvs distance type the distance type -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:250`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:266`_ ### getNNDescentNumIterations @@ -196,7 +210,7 @@ get the number of Iterations to run if building with NN_DESCENT the number of iterations for NN_DESCENT -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:259`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:275`_ ### getHnswHeuristicType @@ -211,7 +225,7 @@ beamWidth. Only consulted under the `Strategy#HEURISTIC` strategy. the `HnswHeuristicType` to hand to cuVS -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:269`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:285`_ ### withWriterThreads @@ -219,8 +233,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWP public Builder withWriterThreads(int writerThreads) ``` -Set the maximum number of threads used for cuVS writes and accelerated-HNSW graph -materialization and serialization. +Set the number of native cuVS writer threads while building the index. Valid range - Minimum: \{@value MIN_WRITER_THREADS\}, Maximum: \{@value MAX_WRITER_THREADS\} Default value - \{@value DEFAULT_WRITER_THREADS\} @@ -234,7 +247,29 @@ Default value - \{@value DEFAULT_WRITER_THREADS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:335`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:353`_ + +### withGraphThreads + +```java +public Builder withGraphThreads(int graphThreads) +``` + +Set the maximum threads per HNSW graph materialization or serialization operation. The count +includes the calling thread. Valid range - Minimum: \{@value MIN_GRAPH_THREADS\}, Maximum: +\{@value MAX_GRAPH_THREADS\}. Default value - \{@value DEFAULT_GRAPH_THREADS\}. + +**Parameters** + +| Name | Description | +| --- | --- | +| `graphThreads` | maximum graph-processing threads per operation | + +**Returns** + +instance of `Builder` + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:366`_ ### withIntermediateGraphDegree @@ -256,7 +291,7 @@ Default value - \{@value DEFAULT_INT_GRAPH_DEGREE\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:348`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:379`_ ### withGraphDegree @@ -278,7 +313,7 @@ Default value - \{@value DEFAULT_GRAPH_DEGREE\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:361`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:392`_ ### withHNSWLayer @@ -300,7 +335,7 @@ Default value - \{@value DEFAULT_HNSW_LAYERS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:374`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:405`_ ### withMaxConn @@ -322,7 +357,7 @@ Default value - \{@value DEFAULT_MAX_CONN\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:387`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:418`_ ### withBeamWidth @@ -344,7 +379,7 @@ Default value - \{@value DEFAULT_BEAM_WIDTH\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:400`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:431`_ ### withCagraGraphBuildAlgo @@ -365,7 +400,7 @@ Default value - NN_DESCENT instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:412`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:443`_ ### withCuVSIvfPqParams @@ -385,7 +420,7 @@ Set the instance of `CuVSIvfPqParams` instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:423`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:454`_ ### withNumMergeWorkers @@ -406,7 +441,7 @@ Default value - \{@value DEFAULT_NUM_MERGE_WORKERS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:435`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:466`_ ### withMergeExecutorService @@ -427,7 +462,7 @@ Default value an instance with one thread instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:447`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:478`_ ### withStrategy @@ -453,7 +488,7 @@ Default value - HEURISTIC instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:464`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:495`_ ### withCuvsDistanceType @@ -473,7 +508,7 @@ Set the CuvsDistanceType instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:475`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:506`_ ### withNNDescentNumIterations @@ -496,7 +531,7 @@ Default value - \{@value DEFAULT_NN_DESCENT_NUM_ITERATIONS\} instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:489`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:520`_ ### withHnswHeuristicType @@ -520,7 +555,7 @@ the equivalent HNSW graph (graph degree = 2 * maxConn). instance of `Builder` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:504`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:535`_ ### build @@ -534,6 +569,6 @@ Create an instance of `AcceleratedHNSWParams` instance of `AcceleratedHNSWParams` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:549`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:582`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:17`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md index 640f8bdaeb..47a7161773 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswutils.md @@ -21,7 +21,7 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens Creates a dummy HNSW graph for a single vector. The graph will have 1 level with 1 node and no neighbors. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:61`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:58`_ ### createMultiLayerHnswGraph @@ -33,7 +33,7 @@ Creates up to `hnswLayers` total layers. Layer 0 uses the full CAGRA graph. Each layer samples `max(2, floor(previousLayerSize / M))` nodes. The value `M` is the ceiling of half the layer-0 graph degree. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:79`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:76`_ ### createMultiLayerHnswGraph @@ -44,7 +44,7 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( int dimensions, CuVSMatrix a Creates a multi-layer HNSW graph from a native matrix without copying the complete dataset to the Java heap. The list view copies only rows selected for an upper layer. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:211`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:239`_ ### writeGraph @@ -71,7 +71,7 @@ a 2D array of offsets | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:360`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:416`_ ### writeMeta @@ -100,7 +100,7 @@ Writes the meta information for the index. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:524`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:589`_ ### printInfoStream @@ -116,7 +116,7 @@ A utility method to print info/debugging messages using InfoStream. | --- | --- | | `msg` | the debugging message to print | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:606`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:671`_ ### writeEmpty @@ -138,7 +138,7 @@ Writes an empty meta information for the field. | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:618`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:683`_ ### quantizeFloatVectorsToBinary @@ -161,7 +161,7 @@ Bits are packed: 8 dimensions per byte. A list of byte binary representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:631`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:696`_ ### quantizeFloatVectorsToScalar @@ -181,6 +181,6 @@ Scalar quantization. A list of byte scalar representation for the input vectors -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:673`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:738`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:38`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java:35`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md index 3367485397..eb0a751516 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-gpubuilthnswgraph.md @@ -31,27 +31,7 @@ Multi-layer constructor that supports arbitrary number of layers. | `layerNodes` | the nodes on the layer | | `layerAdjacencies` | adjacency list | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:55`_ - -### GPUBuiltHnswGraph - -```java -public GPUBuiltHnswGraph( int size, int dimensions, List layerNodes, List layerAdjacencies, int numThreads) throws IOException -``` - -Builds a graph while materializing adjacency rows with the requested number of threads. - -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:61`_ - -### fitsParallelGraphCopyBudget - -```java -static boolean fitsParallelGraphCopyBudget(long rows, long columns) -``` - -Returns whether an INT32 adjacency can be copied without exceeding the native-host budget. - -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:169`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:56`_ ### getNodesOnLevel @@ -61,7 +41,7 @@ public NodesIterator getNodesOnLevel(int level) Get all nodes on a given level as node 0th ordinals. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:256`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:295`_ ### getNeighbors @@ -82,7 +62,7 @@ Get the neighbors for the node and the level it resides. an instance of NeighborArray -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:274`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:313`_ ### seek @@ -92,7 +72,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Move the pointer to exactly the given level's target. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:299`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:338`_ ### nextNeighbor @@ -102,7 +82,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Iterates over the neighbor list. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:309`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:348`_ ### entryNode @@ -112,7 +92,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns graph's entry point on the top level. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:340`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:379`_ ### maxConn @@ -122,7 +102,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap returns M, the maximum number of connections for a node. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:359`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:398`_ ### neighborCount @@ -132,7 +112,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns the neighbor count. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:374`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:413`_ ### size @@ -142,7 +122,7 @@ public int size() Returns the number of nodes in the graph. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:449`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:488`_ ### numLevels @@ -156,7 +136,7 @@ Returns the number of levels in the HNSW graph. the number of levels -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:458`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:497`_ ### dimensions @@ -170,6 +150,6 @@ Gets the vector dimension. the vector dimension -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:467`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:506`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:29`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:27`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md index 18a59572c5..ea991d5f91 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-lucene99acceleratedhnswvectorswriter.md @@ -37,7 +37,7 @@ Initializes `Lucene99AcceleratedHNSWVectorsWriter` | --- | --- | | `IOException` | IOException | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:89`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:90`_ ### addField @@ -47,7 +47,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:131`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:133`_ ### flush @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:223`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:230`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:366`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:377`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:375`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:386`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:395`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:406`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99Accelera Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:405`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:416`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java:55`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md index ae41d679dd..2d058c70f6 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswbinaryquantizedvectorswriter.md @@ -37,7 +37,7 @@ Initializes `LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter` | --- | --- | | `IOException` | IOException | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:79`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:80`_ ### addField @@ -47,7 +47,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:127`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:129`_ ### flush @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:223`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:229`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:309`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:319`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:340`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:350`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:360`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:370`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:369`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:379`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java:57`_ diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md index 8a67d9e920..e4fd6e79ff 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-luceneacceleratedhnswscalarquantizedvectorswriter.md @@ -37,7 +37,7 @@ Initializes `LuceneAcceleratedHNSWScalarQuantizedVectorsWriter` | --- | --- | | `IOException` | IOException | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:89`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:90`_ ### addField @@ -47,7 +47,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Add new field for indexing. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:136`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:138`_ ### flush @@ -57,7 +57,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Build the indexes and writes it to the disk. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:250`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:256`_ ### mergeOneField @@ -67,7 +67,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Write field for merging. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:334`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:344`_ ### finish @@ -77,7 +77,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Called once at the end before close. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:365`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:375`_ ### close @@ -87,7 +87,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:385`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:395`_ ### ramBytesUsed @@ -97,6 +97,6 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAccelerate Returns the memory usage of this object in bytes. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:394`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:404`_ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java:56`_ diff --git a/fern/pages/user_guide/lucene.md b/fern/pages/user_guide/lucene.md index 2f3e49cd20..670a67305f 100644 --- a/fern/pages/user_guide/lucene.md +++ b/fern/pages/user_guide/lucene.md @@ -231,9 +231,7 @@ Both `AcceleratedHNSWParams` and `GPUSearchParams` default to a `HEURISTIC` stra - For the accelerated HNSW codecs, set `maxConn` and `beamWidth`, the HNSW parameters you would tune on the CPU. cuVS derives graph degrees and the build algorithm from them. These two values also configure the CPU fallback writer, so one setting covers both paths. - For the GPU search codec, set `buildQuality`. Higher values spend more build time for a higher-quality graph. This codec also passes `graphDegree` into the heuristic, so leave it at its default unless you intend to cap the graph. -`writerThreads` controls native cuVS writer concurrency for both parameter types. For the accelerated HNSW codecs, it also sets the maximum number of CPU threads used to materialize and serialize the finished graph; the default of one keeps this post-build processing serial. The GPU search codec continues to use the setting only for native writer concurrency and defaults to 32. - -With more than one writer thread, parallel adjacency materialization is eligible for layers with at least 65,536 nodes. A device adjacency is copied once to native host memory when that copy is at most 4 GiB; larger device adjacencies retain serial row access. Level-zero serialization is independently eligible at 65,536 nodes and runs in fixed waves of at most 1,048,576 nodes; upper levels remain serial. +`writerThreads` controls native cuVS build concurrency. The accelerated HNSW codecs default to one writer thread, while the GPU search codec defaults to 32. Accelerated HNSW also offers `graphThreads` for the later CPU graph-materialization and serialization stages; it defaults to one and counts the calling thread. Parallel materialization of a device-backed graph uses a temporary host copy only when the JVM reports enough free physical memory for that copy, the Lucene graph, and an adjacency-sized safety allowance. Otherwise, materialization remains serial. Switching either class to the `CUSTOM` strategy exposes the underlying CAGRA parameters directly, including `graphDegree`, `intermediateGraphDegree`, the graph build algorithm, and its parameters. Use `CUSTOM` only when you have measurements that justify specific values; the defaults derived by cuVS are a better starting point. For background on the parameters themselves, see the [CAGRA indexing guide](/user-guide/api-guides/indexing-guide/cagra) and the [tuning guide](/getting-started/introduction/tuning-indexes). diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java index 458566931f..1f8c13609b 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java @@ -34,6 +34,8 @@ public static enum Strategy { // Bounds for the public CAGRA and HNSW build parameters. public static final int MIN_WRITER_THREADS = 1; public static final int MAX_WRITER_THREADS = 512; + public static final int MIN_GRAPH_THREADS = 1; + public static final int MAX_GRAPH_THREADS = 512; public static final int MIN_INT_GRAPH_DEG = 2; public static final int MAX_INT_GRAPH_DEG = 512; public static final int MIN_GRAPH_DEG = 1; @@ -50,6 +52,7 @@ public static enum Strategy { public static final int MAX_NN_DESCENT_NUM_ITERATIONS = 100; public static final int DEFAULT_WRITER_THREADS = 1; + public static final int DEFAULT_GRAPH_THREADS = 1; public static final int DEFAULT_INT_GRAPH_DEGREE = 128; public static final int DEFAULT_GRAPH_DEGREE = 64; public static final int DEFAULT_HNSW_LAYERS = 1; @@ -75,6 +78,7 @@ public static enum Strategy { }; private final int writerThreads; + private final int graphThreads; private final int intermediateGraphDegree; private final int graphdegree; private final int hnswLayers; @@ -92,8 +96,9 @@ public static enum Strategy { /** * Constructs an instance of {@link AcceleratedHNSWParams} with specific parameter values. * - * @param writerThreads Maximum number of threads to use for cuVS writes and accelerated-HNSW - * graph materialization and serialization. + * @param writerThreads Number of native cuVS writer threads to use. + * @param graphThreads Maximum threads per HNSW graph materialization or serialization operation, + * including the calling thread. * @param intermediateGraphDegree The intermediate graph degree while building the CAGRA index. * @param graphdegree The graph degree to use while building the CAGRA index. * @param hnswLayers The number of HNSW layers to build in the HNSW index. @@ -110,6 +115,7 @@ public static enum Strategy { */ private AcceleratedHNSWParams( int writerThreads, + int graphThreads, int intermediateGraphDegree, int graphdegree, int hnswLayers, @@ -125,6 +131,7 @@ private AcceleratedHNSWParams( HnswHeuristicType hnswHeuristicType) { super(); this.writerThreads = writerThreads; + this.graphThreads = graphThreads; this.intermediateGraphDegree = intermediateGraphDegree; this.graphdegree = graphdegree; this.hnswLayers = hnswLayers; @@ -141,15 +148,24 @@ private AcceleratedHNSWParams( } /** - * Get the maximum thread count for cuVS writes and accelerated-HNSW graph materialization and - * serialization. + * Get the native cuVS writer threads parameter. * - * @return maximum writer and graph-processing thread count + * @return cuVS writer threads parameter */ public int getWriterThreads() { return writerThreads; } + /** + * Get the maximum threads per HNSW graph materialization or serialization operation. The count + * includes the calling thread; shared helper capacity may reduce actual concurrency. + * + * @return HNSW graph processing threads parameter + */ + public int getGraphThreads() { + return graphThreads; + } + /** * Get the intermediate graph degree * @@ -275,6 +291,8 @@ public HnswHeuristicType getHnswHeuristicType() { public String toString() { return "AcceleratedHNSWParams [writerThreads=" + writerThreads + + ", graphThreads=" + + graphThreads + ", intermediateGraphDegree=" + intermediateGraphDegree + ", graphdegree=" @@ -310,6 +328,7 @@ public String toString() { public static class Builder { private int writerThreads = DEFAULT_WRITER_THREADS; + private int graphThreads = DEFAULT_GRAPH_THREADS; private int intermediateGraphDegree = DEFAULT_INT_GRAPH_DEGREE; private int graphdegree = DEFAULT_GRAPH_DEGREE; private int hnswLayers = DEFAULT_HNSW_LAYERS; @@ -325,8 +344,7 @@ public static class Builder { private HnswHeuristicType hnswHeuristicType = DEFAULT_HNSW_HEURISTIC_TYPE; /** - * Set the maximum number of threads used for cuVS writes and accelerated-HNSW graph - * materialization and serialization. + * Set the number of native cuVS writer threads while building the index. * Valid range - Minimum: {@value MIN_WRITER_THREADS}, Maximum: {@value MAX_WRITER_THREADS} * Default value - {@value DEFAULT_WRITER_THREADS} * @@ -338,6 +356,19 @@ public Builder withWriterThreads(int writerThreads) { return this; } + /** + * Set the maximum threads per HNSW graph materialization or serialization operation. The count + * includes the calling thread. Valid range - Minimum: {@value MIN_GRAPH_THREADS}, Maximum: + * {@value MAX_GRAPH_THREADS}. Default value - {@value DEFAULT_GRAPH_THREADS}. + * + * @param graphThreads maximum graph-processing threads per operation + * @return instance of {@link Builder} + */ + public Builder withGraphThreads(int graphThreads) { + this.graphThreads = graphThreads; + return this; + } + /** * Set the intermediate graph degree to use while building CAGRA index * Valid range - Minimum: {@value MIN_INT_GRAPH_DEG}, Maximum: {@value MAX_INT_GRAPH_DEG} @@ -515,6 +546,8 @@ public Builder withHnswHeuristicType(HnswHeuristicType hnswHeuristicType) { private void validate() throws IllegalArgumentException { ParameterValidation.checkRange( "writerThreads", writerThreads, MIN_WRITER_THREADS, MAX_WRITER_THREADS); + ParameterValidation.checkRange( + "graphThreads", graphThreads, MIN_GRAPH_THREADS, MAX_GRAPH_THREADS); ParameterValidation.checkRange( "intermediateGraphDegree", intermediateGraphDegree, MIN_INT_GRAPH_DEG, MAX_INT_GRAPH_DEG); ParameterValidation.checkRange("graphdegree", graphdegree, MIN_GRAPH_DEG, MAX_GRAPH_DEG); @@ -557,6 +590,7 @@ public AcceleratedHNSWParams build() { validate(); return new AcceleratedHNSWParams( writerThreads, + graphThreads, intermediateGraphDegree, graphdegree, hnswLayers, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java index ae5d9b1714..5257af0509 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWUtils.java @@ -21,11 +21,8 @@ import java.util.SortedSet; import java.util.TreeSet; import java.util.concurrent.Callable; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.VectorSimilarityFunction; -import org.apache.lucene.search.TaskExecutor; import org.apache.lucene.store.ByteBuffersDataOutput; import org.apache.lucene.store.DataOutput; import org.apache.lucene.store.IndexOutput; @@ -68,7 +65,7 @@ public static GPUBuiltHnswGraph createSingleVectorHnswGraph(int size, int dimens try (CuVSMatrix adjacencyMatrix = CuVSMatrix.ofArray(singleNodeAdjacency)) { List layerNodes = new ArrayList<>(); layerNodes.add(null); // Layer 0 contains all nodes, so its node list is implicit. - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, List.of(adjacencyMatrix), 1); + return new GPUBuiltHnswGraph(size, dimensions, layerNodes, List.of(adjacencyMatrix)); } } @@ -88,7 +85,14 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( QuantizationType quantization) throws Throwable { return createMultiLayerHnswGraph( - size, dimensions, adjacencyListMatrix, vectors, hnswLayers, params, quantization, 1); + size, + dimensions, + adjacencyListMatrix, + vectors, + hnswLayers, + params, + quantization, + AcceleratedHNSWParams.DEFAULT_GRAPH_THREADS); } private static GPUBuiltHnswGraph createMultiLayerHnswGraph( @@ -99,7 +103,30 @@ private static GPUBuiltHnswGraph createMultiLayerHnswGraph( int hnswLayers, CagraIndexParams params, QuantizationType quantization, - int numThreads) + int graphThreads) + throws Throwable { + return createMultiLayerHnswGraph( + size, + dimensions, + adjacencyListMatrix, + vectors, + hnswLayers, + params, + quantization, + graphThreads, + GraphProcessingTrace.disabled()); + } + + private static GPUBuiltHnswGraph createMultiLayerHnswGraph( + int size, + int dimensions, + CuVSMatrix adjacencyListMatrix, + List vectors, + int hnswLayers, + CagraIndexParams params, + QuantizationType quantization, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) throws Throwable { int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); @@ -188,7 +215,8 @@ private static GPUBuiltHnswGraph createMultiLayerHnswGraph( } // The graph eagerly copies all adjacency rows, so generated upper matrices can now close. - return new GPUBuiltHnswGraph(size, dimensions, layerNodes, layerAdjacencies, numThreads); + return new GPUBuiltHnswGraph( + size, dimensions, layerNodes, layerAdjacencies, graphThreads, graphProcessingTrace); } catch (Throwable t) { failure = t; throw t; @@ -218,7 +246,13 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( QuantizationType quantization) throws Throwable { return createMultiLayerHnswGraph( - dimensions, adjacencyListMatrix, vectorDataset, hnswLayers, params, quantization, 1); + dimensions, + adjacencyListMatrix, + vectorDataset, + hnswLayers, + params, + quantization, + AcceleratedHNSWParams.DEFAULT_GRAPH_THREADS); } static GPUBuiltHnswGraph createMultiLayerHnswGraph( @@ -228,7 +262,28 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( int hnswLayers, CagraIndexParams params, QuantizationType quantization, - int numThreads) + int graphThreads) + throws Throwable { + return createMultiLayerHnswGraph( + dimensions, + adjacencyListMatrix, + vectorDataset, + hnswLayers, + params, + quantization, + graphThreads, + GraphProcessingTrace.disabled()); + } + + static GPUBuiltHnswGraph createMultiLayerHnswGraph( + int dimensions, + CuVSMatrix adjacencyListMatrix, + CuVSMatrix vectorDataset, + int hnswLayers, + CagraIndexParams params, + QuantizationType quantization, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) throws Throwable { int size = Math.toIntExact(vectorDataset.size()); // Matrix columns are the stored width: binary vectors are bit-packed, while scalar and float @@ -262,7 +317,8 @@ public int size() { hnswLayers, params, quantization, - numThreads); + graphThreads, + graphProcessingTrace); } private static Throwable closeUpperLayerAdjacencies(List layerAdjacencies) { @@ -360,29 +416,46 @@ private static CuVSMatrix buildCagraGraphForSubset( */ public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) throws IOException { - return writeGraph(graph, vectorIndex, 1); + return writeGraph(graph, vectorIndex, AcceleratedHNSWParams.DEFAULT_GRAPH_THREADS); + } + + static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int graphThreads) + throws IOException { + return writeGraph(graph, vectorIndex, graphThreads, GraphProcessingTrace.disabled()); } - public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex, int numThreads) + static int[][] writeGraph( + GPUBuiltHnswGraph graph, + IndexOutput vectorIndex, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) throws IOException { int countOnLevel0 = graph.size(); int numLevels = graph.numLevels(); int[][] offsets = new int[numLevels][]; - - // Each level-0 node's delta/VInt block is independent. Encode those blocks in parallel in fixed - // waves, then concatenate thread buffers in node order. Higher levels are much smaller and stay - // serial. Both paths write the same blocks and offsets in the same order. - // graph.maxConn() scans every layer-0 adjacency row (O(graph size)); compute it once here - // rather than per level/per task below. int maxConn = graph.maxConn(); int[] level0Nodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(0)); offsets[0] = new int[level0Nodes.length]; - if (numThreads > 1 && level0Nodes.length >= PARALLEL_MIN_NODES) { + if (graphThreads > 1 && level0Nodes.length >= GPUBuiltHnswGraph.PARALLEL_MIN_NODES) { writeLevel0Parallel( - graph, vectorIndex, level0Nodes, offsets[0], countOnLevel0, maxConn, numThreads); + graph, vectorIndex, level0Nodes, offsets[0], countOnLevel0, maxConn, graphThreads); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.SERIALIZATION, + GraphProcessingTrace.Mode.PARALLEL, + GraphProcessingTrace.Reason.ABOVE_THRESHOLD, + graphThreads, + level0Nodes.length); } else { writeLevelSerial(graph, vectorIndex, 0, level0Nodes, offsets[0], countOnLevel0, maxConn); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.SERIALIZATION, + GraphProcessingTrace.Mode.SERIAL, + graphThreads <= 1 + ? GraphProcessingTrace.Reason.SINGLE_THREAD + : GraphProcessingTrace.Reason.BELOW_THRESHOLD, + graphThreads, + level0Nodes.length); } for (int level = 1; level < numLevels; level++) { @@ -394,36 +467,36 @@ public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorInde return offsets; } - /** Node count below which parallel level-0 serialization is not worth the overhead. */ - static final int PARALLEL_MIN_NODES = 1 << 16; + /** + * Fixed operational guardrail on nodes processed before task-local buffers are concatenated when + * the encoded-byte limit would otherwise permit a very large wave. Unlike the byte limit below, + * this is a policy cap rather than an encoded-size calculation. + */ + static final int MAX_SERIALIZATION_WAVE_NODES = 1 << 20; - /** Nodes per wave, bounding the number of nodes buffered independently of dataset size. */ - static final int SERIALIZATION_WAVE_NODES = 1 << 20; + /** Maximum worst-case encoded payload buffered by one serialization wave. */ + static final long MAX_SERIALIZED_BYTES_PER_WAVE = 64L << 20; + + private static final int MAX_VINT_BYTES = 5; - /** Serially encodes a level's nodes into {@code out}, recording per-node byte lengths. */ private static void writeLevelSerial( GPUBuiltHnswGraph graph, IndexOutput out, int level, - int[] sortedNodes, + int[] nodes, int[] offsets, int countOnLevel0, int maxConn) throws IOException { int[] scratch = new int[maxConn * 2]; - int idx = 0; - for (int node : sortedNodes) { + for (int i = 0; i < nodes.length; i++) { long start = out.getFilePointer(); - encodeNode(graph.getNeighbors(level, node), scratch, out, countOnLevel0); - offsets[idx++] = Math.toIntExact(out.getFilePointer() - start); + encodeNode(graph.getNeighbors(level, nodes[i]), scratch, out, countOnLevel0); + offsets[i] = Math.toIntExact(out.getFilePointer() - start); } } - /** - * Encodes level 0 in parallel: within fixed-size waves, threads encode contiguous node - * sub-ranges into per-thread buffers, which are then concatenated to {@code out} in node order - * (identical layout to the serial path). - */ + /** Encodes level zero in bounded waves, then concatenates buffers in node order. */ private static void writeLevel0Parallel( GPUBuiltHnswGraph graph, IndexOutput out, @@ -431,76 +504,68 @@ private static void writeLevel0Parallel( int[] offsets, int countOnLevel0, int maxConn, - int numThreads) + int graphThreads) throws IOException { - // invokeAll joins every task before it returns, including tasks still running when a sibling - // fails, so `buffers` is never concatenated while a worker might still be writing into it. - // It also runs one share of each wave on the calling thread instead of parking it, so the - // pool only has to cover the other ranges. - ExecutorService pool = Executors.newFixedThreadPool(Math.max(1, numThreads - 1)); - try { - TaskExecutor executor = new TaskExecutor(pool); - int n = nodes.length; - for (int waveStart = 0; waveStart < n; ) { - int waveEnd = (int) Math.min(n, (long) waveStart + SERIALIZATION_WAVE_NODES); - int perThread = (waveEnd - waveStart + numThreads - 1) / numThreads; - - ByteBuffersDataOutput[] buffers = new ByteBuffersDataOutput[numThreads]; - List> tasks = new ArrayList<>(numThreads); - for (int t = 0; t < numThreads; t++) { - final int subStart = waveStart + t * perThread; - final int subEnd = Math.min(subStart + perThread, waveEnd); - final int slot = t; - if (subStart >= subEnd) { - continue; - } - tasks.add( - () -> { - ByteBuffersDataOutput buffer = new ByteBuffersDataOutput(); - int[] scratch = new int[maxConn * 2]; - for (int i = subStart; i < subEnd; i++) { - long before = buffer.size(); - encodeNode(graph.getNeighbors(0, nodes[i]), scratch, buffer, countOnLevel0); - offsets[i] = Math.toIntExact(buffer.size() - before); - } - buffers[slot] = buffer; - return null; - }); + int waveNodes = serializationWaveNodes(maxConn); + for (int waveStart = 0; waveStart < nodes.length; ) { + int waveEnd = (int) Math.min(nodes.length, (long) waveStart + waveNodes); + int nodesPerTask = Math.ceilDiv(waveEnd - waveStart, graphThreads); + ByteBuffersDataOutput[] buffers = new ByteBuffersDataOutput[graphThreads]; + List> tasks = new ArrayList<>(graphThreads); + for (int task = 0; task < graphThreads; task++) { + int start = waveStart + task * nodesPerTask; + int end = Math.min(start + nodesPerTask, waveEnd); + int bufferIndex = task; + if (start >= end) { + break; } - executor.invokeAll(tasks); - // Concatenate in thread order (== node order), preserving the serial byte layout. - for (ByteBuffersDataOutput buffer : buffers) { - if (buffer != null) { - buffer.copyTo(out); - } + tasks.add( + () -> { + ByteBuffersDataOutput buffer = new ByteBuffersDataOutput(); + int[] scratch = new int[maxConn * 2]; + for (int i = start; i < end; i++) { + long before = buffer.size(); + encodeNode(graph.getNeighbors(0, nodes[i]), scratch, buffer, countOnLevel0); + offsets[i] = Math.toIntExact(buffer.size() - before); + } + buffers[bufferIndex] = buffer; + return null; + }); + } + GraphWorkExecutor.invokeAll(tasks); + for (ByteBuffersDataOutput buffer : buffers) { + if (buffer != null) { + buffer.copyTo(out); } - waveStart = waveEnd; } - } finally { - pool.shutdown(); + waveStart = waveEnd; } } - /** - * Sorts, delta-encodes and de-duplicates a node's neighbors and writes the block (VInt size + VInt - * deltas) to {@code out}. Shared by the serial and parallel paths so encoding is identical. - */ + static int serializationWaveNodes(int maxConn) { + if (maxConn < 0) { + throw new IllegalArgumentException("maxConn must not be negative"); + } + long maxBytesPerNode = Math.addExact(MAX_VINT_BYTES, (long) maxConn * MAX_VINT_BYTES); + long byteBoundedNodes = Math.max(1, MAX_SERIALIZED_BYTES_PER_WAVE / maxBytesPerNode); + return (int) Math.min(MAX_SERIALIZATION_WAVE_NODES, byteBoundedNodes); + } + private static void encodeNode( NeighborArray neighbors, int[] scratch, DataOutput out, int countOnLevel0) throws IOException { - int size = neighbors == null ? 0 : neighbors.size(); + int size = neighbors.size(); int actualSize = 0; if (size > 0) { - int[] nnodes = neighbors.nodes(); - Arrays.sort(nnodes, 0, size); - scratch[0] = nnodes[0]; + int[] nodes = neighbors.nodes(); + Arrays.sort(nodes, 0, size); + scratch[0] = nodes[0]; actualSize = 1; for (int i = 1; i < size; i++) { - assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; - if (nnodes[i - 1] == nnodes[i]) { - continue; + assert nodes[i] < countOnLevel0 : "node too large: " + nodes[i] + ">=" + countOnLevel0; + if (nodes[i - 1] != nodes[i]) { + scratch[actualSize++] = nodes[i] - nodes[i - 1]; } - scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; } } out.writeVInt(actualSize); diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java index f6f15eef5f..61c73b10ab 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java @@ -13,11 +13,9 @@ import java.io.IOException; import java.util.ArrayList; import java.util.List; +import java.util.Optional; import java.util.concurrent.Callable; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import java.util.function.Supplier; -import org.apache.lucene.search.TaskExecutor; import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.NeighborArray; @@ -39,6 +37,9 @@ public class GPUBuiltHnswGraph extends HnswGraph { // Layer 0 is special - it contains all nodes private final NeighborArray[] layer0Neighbors; + /** Node count below which parallel graph processing is not worth its overhead. */ + static final int PARALLEL_MIN_NODES = 1 << 16; + private record MaterializedGraph( int numLevels, List layerNodes, @@ -58,15 +59,33 @@ public GPUBuiltHnswGraph( this(size, dimensions, materializeSerial(size, layerNodes, layerAdjacencies)); } - /** Builds a graph while materializing adjacency rows with the requested number of threads. */ - public GPUBuiltHnswGraph( + // Builds a graph while materializing adjacency rows with up to graphThreads threads. + GPUBuiltHnswGraph( int size, int dimensions, List layerNodes, List layerAdjacencies, - int numThreads) + int graphThreads) throws IOException { - this(size, dimensions, materialize(size, layerNodes, layerAdjacencies, numThreads)); + this( + size, + dimensions, + materialize( + size, layerNodes, layerAdjacencies, graphThreads, GraphProcessingTrace.disabled())); + } + + GPUBuiltHnswGraph( + int size, + int dimensions, + List layerNodes, + List layerAdjacencies, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) + throws IOException { + this( + size, + dimensions, + materialize(size, layerNodes, layerAdjacencies, graphThreads, graphProcessingTrace)); } private GPUBuiltHnswGraph(int size, int dimensions, MaterializedGraph graph) { @@ -80,107 +99,136 @@ private GPUBuiltHnswGraph(int size, int dimensions, MaterializedGraph graph) { private static MaterializedGraph materializeSerial( int size, List layerNodes, List layerAdjacencies) { - List upperLayerNodes = new ArrayList<>(); - List upperLayerNeighbors = new ArrayList<>(); - NeighborArray[] baseLayerNeighbors = fillNeighborArraySerial(layerAdjacencies.get(0), size); - - for (int level = 1; level < layerAdjacencies.size(); level++) { - int[] nodes = layerNodes.get(level); - upperLayerNodes.add(nodes); - upperLayerNeighbors.add(fillNeighborArraySerial(layerAdjacencies.get(level), nodes.length)); + try { + return materialize(size, layerNodes, layerAdjacencies, 1, GraphProcessingTrace.disabled()); + } catch (IOException impossible) { + throw new AssertionError( + "serial graph materialization cannot fail with IOException", impossible); } - return new MaterializedGraph( - layerAdjacencies.size(), upperLayerNodes, baseLayerNeighbors, upperLayerNeighbors); } private static MaterializedGraph materialize( - int size, List layerNodes, List layerAdjacencies, int numThreads) + int size, + List layerNodes, + List layerAdjacencies, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) throws IOException { - if (numThreads <= 1) { - return materializeSerial(size, layerNodes, layerAdjacencies); - } - List upperLayerNodes = new ArrayList<>(); List upperLayerNeighbors = new ArrayList<>(); NeighborArray[] baseLayerNeighbors = - fillNeighborArray(layerAdjacencies.get(0), size, numThreads); + fillNeighborArray(layerAdjacencies.get(0), size, graphThreads, graphProcessingTrace); for (int level = 1; level < layerAdjacencies.size(); level++) { int[] nodes = layerNodes.get(level); upperLayerNodes.add(nodes); upperLayerNeighbors.add( - fillNeighborArray(layerAdjacencies.get(level), nodes.length, numThreads)); + fillNeighborArray( + layerAdjacencies.get(level), nodes.length, graphThreads, graphProcessingTrace)); } return new MaterializedGraph( layerAdjacencies.size(), upperLayerNodes, baseLayerNeighbors, upperLayerNeighbors); } - /** Node count below which parallel materialization is not worth the thread overhead. */ - static final int PARALLEL_MIN_NODES = 1 << 16; - - /** - * Maximum temporary native-host copy used to make a device adjacency safe for concurrent reads. - * Larger device matrices retain serial row access instead of risking a full-matrix native-memory - * spike on top of the Java {@link NeighborArray} representation. - */ - static final long MAX_PARALLEL_GRAPH_COPY_BYTES = 4L << 30; - /** - * Materializes the adjacency matrix into on-heap {@link NeighborArray}s, one per node. - * - *

The serial path reads the adjacency directly (a device matrix's {@code getRow} is safe - * single-threaded). The parallel path cannot: the CAGRA layer-0 adjacency is a device matrix whose - * {@code getRow} uses a shared, stateful buffered reader that is not safe for concurrent access, so - * it is pulled to host once (a single bulk device->host copy) before materializing disjoint node - * ranges concurrently. Host matrices (the upper layers, built via {@link CuVSMatrix#ofArray}) are - * read directly in both paths. + * Materializes an adjacency matrix into heap-backed neighbor arrays. Device matrices are copied + * to host before parallel reads because their row reader is stateful. * * @param adjacency instance of adjacency CuVSMatrix * @param size the number of nodes - * @param numThreads threads to use (1, or fewer than {@value #PARALLEL_MIN_NODES} nodes = serial) + * @param graphThreads maximum threads for this operation, including the caller * @return the NeighborArray */ - private static NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size, int numThreads) + private static NeighborArray[] fillNeighborArray( + CuVSMatrix adjacency, int size, int graphThreads, GraphProcessingTrace graphProcessingTrace) throws IOException { - NeighborArray[] neighbors = new NeighborArray[size]; - if (numThreads <= 1 - || size < PARALLEL_MIN_NODES - || (adjacency instanceof CuVSDeviceMatrix - && !fitsParallelGraphCopyBudget(adjacency.size(), adjacency.columns()))) { - fillNeighborRange(adjacency, neighbors, 0, size); + if (graphThreads <= 1 || size < PARALLEL_MIN_NODES) { + NeighborArray[] neighbors = fillNeighborArraySerial(adjacency, size); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.MATERIALIZATION, + GraphProcessingTrace.Mode.SERIAL, + graphThreads <= 1 + ? GraphProcessingTrace.Reason.SINGLE_THREAD + : GraphProcessingTrace.Reason.BELOW_THRESHOLD, + graphThreads, + size); return neighbors; } + if (adjacency instanceof CuVSDeviceMatrix deviceAdjacency) { - try (CuVSHostMatrix hostCopy = copyToHost(deviceAdjacency)) { - fillNeighborArrayParallel(hostCopy, neighbors, size, numThreads); - } - return neighbors; + return materializeDeviceAdjacency( + deviceAdjacency, + size, + graphThreads, + GraphCopyMemoryBudget.system(), + () -> newHostMatrix(deviceAdjacency), + graphProcessingTrace); } - fillNeighborArrayParallel(adjacency, neighbors, size, numThreads); + + NeighborArray[] neighbors = fillNeighborArrayParallel(adjacency, size, graphThreads); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.MATERIALIZATION, + GraphProcessingTrace.Mode.PARALLEL, + GraphProcessingTrace.Reason.HOST_SOURCE, + graphThreads, + size); return neighbors; } - private static NeighborArray[] fillNeighborArraySerial(CuVSMatrix adjacency, int size) { + private static NeighborArray[] fillNeighborArraySerial(CuVSMatrix source, int size) { NeighborArray[] neighbors = new NeighborArray[size]; - fillNeighborRange(adjacency, neighbors, 0, size); + fillNeighborRange(source, neighbors, 0, size); return neighbors; } - /** Returns whether an INT32 adjacency can be copied without exceeding the native-host budget. */ - static boolean fitsParallelGraphCopyBudget(long rows, long columns) { - if (rows < 0 || columns < 0) { - return false; + static NeighborArray[] materializeDeviceAdjacency( + CuVSDeviceMatrix source, + int size, + int graphThreads, + GraphCopyMemoryBudget memoryBudget, + Supplier hostCopyFactory) + throws IOException { + return materializeDeviceAdjacency( + source, size, graphThreads, memoryBudget, hostCopyFactory, GraphProcessingTrace.disabled()); + } + + private static NeighborArray[] materializeDeviceAdjacency( + CuVSDeviceMatrix source, + int size, + int graphThreads, + GraphCopyMemoryBudget memoryBudget, + Supplier hostCopyFactory, + GraphProcessingTrace graphProcessingTrace) + throws IOException { + Optional reservation = + memoryBudget.tryReserve(source.size(), source.columns()); + if (reservation.isEmpty()) { + NeighborArray[] neighbors = fillNeighborArraySerial(source, size); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.MATERIALIZATION, + GraphProcessingTrace.Mode.SERIAL, + GraphProcessingTrace.Reason.MEMORY_ADMISSION_DENIED, + graphThreads, + size); + return neighbors; } - if (rows == 0 || columns == 0) { - return true; + try (GraphCopyMemoryBudget.Reservation ignored = reservation.orElseThrow(); + CuVSHostMatrix hostCopy = copyToHost(source, hostCopyFactory)) { + NeighborArray[] neighbors = fillNeighborArrayParallel(hostCopy, size, graphThreads); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.MATERIALIZATION, + GraphProcessingTrace.Mode.PARALLEL, + GraphProcessingTrace.Reason.DEVICE_HOST_COPY, + graphThreads, + size); + return neighbors; } - return rows <= MAX_PARALLEL_GRAPH_COPY_BYTES / Integer.BYTES / columns; } - private static CuVSHostMatrix copyToHost(CuVSDeviceMatrix source) { + private static CuVSHostMatrix newHostMatrix(CuVSDeviceMatrix source) { try (CuVSMatrix.Builder builder = CuVSMatrix.hostBuilder(source.size(), source.columns(), source.dataType())) { - return copyToHost(source, builder::build); + return builder.build(); } } @@ -202,49 +250,40 @@ static CuVSHostMatrix copyToHost( } } - /** - * Materializes disjoint node ranges concurrently. Each thread writes its own slots of {@code - * neighbors} and its own {@link NeighborArray} instances, so no synchronization is needed; {@code - * source} must be a host matrix (stateless {@code getRow}). - */ - private static void fillNeighborArrayParallel( - CuVSMatrix source, NeighborArray[] neighbors, int size, int numThreads) throws IOException { - ExecutorService pool = Executors.newFixedThreadPool(Math.max(1, numThreads - 1)); - try { - int perThread = (size + numThreads - 1) / numThreads; - List> tasks = new ArrayList<>(numThreads); - for (int t = 0; t < numThreads; t++) { - final int start = t * perThread; - final int end = Math.min(start + perThread, size); - if (start >= end) { - break; - } - tasks.add( - () -> { - fillNeighborRange(source, neighbors, start, end); - return null; - }); + private static NeighborArray[] fillNeighborArrayParallel( + CuVSMatrix source, int size, int graphThreads) throws IOException { + NeighborArray[] neighbors = new NeighborArray[size]; + int nodesPerTask = Math.ceilDiv(size, graphThreads); + List> tasks = new ArrayList<>(graphThreads); + for (int task = 0; task < graphThreads; task++) { + int start = task * nodesPerTask; + int end = Math.min(start + nodesPerTask, size); + if (start >= end) { + break; } - new TaskExecutor(pool).invokeAll(tasks); - } finally { - pool.shutdown(); + tasks.add( + () -> { + fillNeighborRange(source, neighbors, start, end); + return null; + }); } + GraphWorkExecutor.invokeAll(tasks); + return neighbors; } - /** Fills {@code neighbors[start, end)} from the adjacency rows. */ private static void fillNeighborRange( CuVSMatrix source, NeighborArray[] neighbors, int start, int end) { for (int i = start; i < end; i++) { RowView rv = source.getRow(i); if (rv != null && rv.size() > 0) { - NeighborArray na = new NeighborArray((int) rv.size(), true); + NeighborArray nodeNeighbors = new NeighborArray((int) rv.size(), true); for (int j = 0; j < rv.size(); j++) { int neighbor = rv.getAsInt(j); if (neighbor >= 0) { - na.addInOrder(neighbor, 1.0f - (j * 0.001f)); + nodeNeighbors.addInOrder(neighbor, 1.0f - (j * 0.001f)); } } - neighbors[i] = na; + neighbors[i] = nodeNeighbors; } else { neighbors[i] = new NeighborArray(0, true); } diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java new file mode 100644 index 0000000000..0a03e884d8 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java @@ -0,0 +1,137 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.sun.management.OperatingSystemMXBean; +import java.lang.management.ManagementFactory; +import java.util.Objects; +import java.util.Optional; +import org.apache.lucene.util.RamUsageEstimator; +import org.apache.lucene.util.hnsw.NeighborArray; + +/** Coordinates temporary native graph copies across concurrent segment flushes. */ +final class GraphCopyMemoryBudget { + private static final long NEIGHBOR_ARRAY_SHALLOW_BYTES = + RamUsageEstimator.shallowSizeOfInstance(NeighborArray.class); + + private static final GraphCopyMemoryBudget SYSTEM = + new GraphCopyMemoryBudget(GraphCopyMemoryBudget::readSystemMemory); + + private final MemoryProbe memoryProbe; + private long reservedHeadroomBytes; + + GraphCopyMemoryBudget(MemoryProbe memoryProbe) { + this.memoryProbe = Objects.requireNonNull(memoryProbe); + } + + static GraphCopyMemoryBudget system() { + return SYSTEM; + } + + /** + * Tries to reserve enough observed free memory for one graph copy and its materialized graph. + * Callers in the same class loader share reservations. This is cooperative admission control, + * not an operating-system memory guarantee. + */ + synchronized Optional tryReserve(long rows, long columns) { + long requiredHeadroom = requiredHeadroom(rows, columns); + if (requiredHeadroom < 0) { + return Optional.empty(); + } + + MemorySnapshot memory; + try { + memory = memoryProbe.read(); + } catch (RuntimeException unavailable) { + return Optional.empty(); + } + if (memory == null + || memory.totalBytes() <= 0 + || memory.freeBytes() < 0 + || memory.freeBytes() > memory.totalBytes()) { + return Optional.empty(); + } + + if (reservedHeadroomBytes > memory.freeBytes() + || requiredHeadroom > memory.freeBytes() - reservedHeadroomBytes) { + return Optional.empty(); + } + reservedHeadroomBytes += requiredHeadroom; + return Optional.of(new Reservation(this, requiredHeadroom)); + } + + /** + * Estimates peak allocation from the actual matrix shape and current JVM object layout. Besides + * the native host copy and materialized Lucene graph, one adjacency-sized allowance protects + * against allocation races and estimation error while the copy is in flight. + */ + static long requiredHeadroom(long rows, long columns) { + if (rows <= 0 || rows > Integer.MAX_VALUE || columns <= 0 || columns > Integer.MAX_VALUE) { + return -1; + } + try { + long adjacencyBytes = Math.multiplyExact(Math.multiplyExact(rows, columns), Integer.BYTES); + long neighborReferences = arraySize(rows, RamUsageEstimator.NUM_BYTES_OBJECT_REF); + long nodeIds = arraySize(columns, Integer.BYTES); + long scores = arraySize(columns, Float.BYTES); + long bytesPerNode = Math.addExact(NEIGHBOR_ARRAY_SHALLOW_BYTES, nodeIds); + bytesPerNode = Math.addExact(bytesPerNode, scores); + long luceneGraphBytes = + Math.addExact(neighborReferences, Math.multiplyExact(rows, bytesPerNode)); + return Math.addExact(Math.multiplyExact(adjacencyBytes, 2), luceneGraphBytes); + } catch (ArithmeticException overflow) { + return -1; + } + } + + private static long arraySize(long length, int bytesPerElement) { + long unaligned = + Math.addExact( + RamUsageEstimator.NUM_BYTES_ARRAY_HEADER, Math.multiplyExact(length, bytesPerElement)); + long alignment = RamUsageEstimator.NUM_BYTES_OBJECT_ALIGNMENT; + long remainder = unaligned % alignment; + return remainder == 0 ? unaligned : Math.addExact(unaligned, alignment - remainder); + } + + private synchronized void release(Reservation reservation) { + if (reservation.released) { + return; + } + reservedHeadroomBytes -= reservation.headroomBytes; + reservation.released = true; + } + + private static MemorySnapshot readSystemMemory() { + java.lang.management.OperatingSystemMXBean platformBean = + ManagementFactory.getOperatingSystemMXBean(); + if (platformBean instanceof OperatingSystemMXBean osBean) { + return new MemorySnapshot(osBean.getTotalMemorySize(), osBean.getFreeMemorySize()); + } + return null; + } + + @FunctionalInterface + interface MemoryProbe { + MemorySnapshot read(); + } + + record MemorySnapshot(long totalBytes, long freeBytes) {} + + static final class Reservation implements AutoCloseable { + private final GraphCopyMemoryBudget budget; + private final long headroomBytes; + private boolean released; + + private Reservation(GraphCopyMemoryBudget budget, long headroomBytes) { + this.budget = budget; + this.headroomBytes = headroomBytes; + } + + @Override + public void close() { + budget.release(this); + } + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphProcessingTrace.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphProcessingTrace.java new file mode 100644 index 0000000000..54f271fd07 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphProcessingTrace.java @@ -0,0 +1,71 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.util.Locale; +import org.apache.lucene.util.InfoStream; + +/** Reports the graph-processing path selected for one Lucene writer. */ +final class GraphProcessingTrace { + + enum Stage { + MATERIALIZATION, + SERIALIZATION + } + + enum Mode { + SERIAL, + PARALLEL + } + + enum Reason { + ABOVE_THRESHOLD, + BELOW_THRESHOLD, + DEVICE_HOST_COPY, + HOST_SOURCE, + MEMORY_ADMISSION_DENIED, + SINGLE_THREAD + } + + private static final GraphProcessingTrace DISABLED = new GraphProcessingTrace(null, null); + + private final InfoStream infoStream; + private final String component; + + private GraphProcessingTrace(InfoStream infoStream, String component) { + this.infoStream = infoStream; + this.component = component; + } + + static GraphProcessingTrace disabled() { + return DISABLED; + } + + static GraphProcessingTrace toInfoStream(InfoStream infoStream, String component) { + return new GraphProcessingTrace(infoStream, component); + } + + void record(Stage stage, Mode mode, Reason reason, int requestedThreads, int nodes) { + if (infoStream == null || !infoStream.isEnabled(component)) { + return; + } + infoStream.message( + component, + "graph-processing stage=" + + lowerCase(stage) + + " mode=" + + lowerCase(mode) + + " reason=" + + lowerCase(reason) + + " requestedThreads=" + + requestedThreads + + " nodes=" + + nodes); + } + + private static String lowerCase(Enum value) { + return value.name().toLowerCase(Locale.ROOT).replace('_', '-'); + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphWorkExecutor.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphWorkExecutor.java new file mode 100644 index 0000000000..cb715f1fd8 --- /dev/null +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphWorkExecutor.java @@ -0,0 +1,167 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.FutureTask; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.SynchronousQueue; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.lucene.util.IOUtils; +import org.apache.lucene.util.ThreadInterruptedException; + +/** Shares a bounded set of helper threads across accelerated-HNSW graph operations. */ +final class GraphWorkExecutor { + static final String THREAD_NAME_PREFIX = "cuvs-hnsw-graph-worker-"; + + private static final int MAX_WORKERS = + defaultMaxWorkers(Runtime.getRuntime().availableProcessors()); + private static final long KEEP_ALIVE_SECONDS = 1; + private static final AtomicInteger NEXT_THREAD_ID = new AtomicInteger(); + private static final ThreadPoolExecutor EXECUTOR = + newExecutor(MAX_WORKERS, KEEP_ALIVE_SECONDS, TimeUnit.SECONDS); + + private GraphWorkExecutor() {} + + static void invokeAll(List> tasks) throws IOException { + invokeAll(EXECUTOR, tasks); + } + + /** + * Runs one task on the calling thread and offers the others to {@code executor}. Direct handoff + * and caller execution apply backpressure without retaining graph work in a queue. + */ + static void invokeAll(Executor executor, List> tasks) throws IOException { + Objects.requireNonNull(executor, "executor"); + Objects.requireNonNull(tasks, "tasks"); + if (tasks.isEmpty()) { + return; + } + AtomicBoolean taskFailed = new AtomicBoolean(); + List> preparedTasks = new ArrayList<>(tasks.size()); + for (Callable task : tasks) { + preparedTasks.add( + new FutureTask<>(stopAfterFailure(Objects.requireNonNull(task, "task"), taskFailed))); + } + + List> acceptedTasks = new ArrayList<>(preparedTasks.size()); + Throwable failure = null; + for (int i = 0; i < preparedTasks.size(); i++) { + FutureTask task = preparedTasks.get(i); + if (i == preparedTasks.size() - 1) { + acceptedTasks.add(task); + task.run(); + break; + } + try { + executor.execute(task); + acceptedTasks.add(task); + } catch (RejectedExecutionException rejected) { + acceptedTasks.add(task); + task.run(); + } catch (RuntimeException | Error submissionFailure) { + taskFailed.set(true); + failure = submissionFailure; + break; + } + } + + failure = collectFailures(acceptedTasks, failure); + if (failure != null) { + throw IOUtils.rethrowAlways(failure); + } + } + + static ThreadPoolExecutor newExecutor(int maxWorkers, long keepAlive, TimeUnit unit) { + if (maxWorkers < 1) { + throw new IllegalArgumentException("maxWorkers must be positive"); + } + if (keepAlive <= 0) { + throw new IllegalArgumentException("keepAlive must be positive"); + } + Objects.requireNonNull(unit, "unit"); + + return new ThreadPoolExecutor( + 0, + maxWorkers, + keepAlive, + unit, + new SynchronousQueue<>(), + workerThreadFactory(), + new ThreadPoolExecutor.AbortPolicy()); + } + + static int defaultMaxWorkers(int availableProcessors) { + return availableProcessors <= 1 ? 1 : availableProcessors - 1; + } + + private static Callable stopAfterFailure(Callable task, AtomicBoolean taskFailed) { + return () -> { + if (taskFailed.get()) { + return null; + } + try { + return task.call(); + } catch (Exception | Error failure) { + taskFailed.set(true); + throw failure; + } + }; + } + + private static ThreadFactory workerThreadFactory() { + return task -> { + Thread worker = + new Thread(null, task, THREAD_NAME_PREFIX + NEXT_THREAD_ID.incrementAndGet(), 0, false); + worker.setDaemon(true); + worker.setPriority(Thread.NORM_PRIORITY); + worker.setContextClassLoader(null); + return worker; + }; + } + + /** Waits for all accepted work before returning, even after failure or interruption. */ + private static Throwable collectFailures(List> tasks, Throwable failure) { + boolean interrupted = false; + for (FutureTask task : tasks) { + while (true) { + try { + task.get(); + break; + } catch (InterruptedException interruption) { + interrupted = true; + failure = addFailure(failure, new ThreadInterruptedException(interruption)); + } catch (ExecutionException taskFailure) { + failure = addFailure(failure, taskFailure.getCause()); + break; + } + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + return failure; + } + + private static Throwable addFailure(Throwable failure, Throwable next) { + if (failure == null) { + return next; + } + if (failure != next) { + failure.addSuppressed(next); + } + return failure; + } +} diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java index 8a27faf957..8f52eae6e0 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/Lucene99AcceleratedHNSWVectorsWriter.java @@ -64,6 +64,7 @@ public class Lucene99AcceleratedHNSWVectorsWriter extends KnnVectorsWriter { private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; + private final GraphProcessingTrace graphProcessingTrace; private IndexOutput hnswMeta = null; private IndexOutput hnswVectorIndex = null; private String vemFileName; @@ -95,6 +96,7 @@ public Lucene99AcceleratedHNSWVectorsWriter( super(); this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; + this.graphProcessingTrace = GraphProcessingTrace.toInfoStream(infoStream, COMPONENT); this.acceleratedHNSWParams = acceleratedHNSWParams; vemFileName = IndexFileNames.segmentFileName( @@ -186,10 +188,15 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.NONE, - acceleratedHNSWParams.getWriterThreads()); + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -275,7 +282,11 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java index 993fb8a58b..4fa02eb1d5 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter.java @@ -63,6 +63,7 @@ public class LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter extends KnnVector private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; + private final GraphProcessingTrace graphProcessingTrace; private final AcceleratedHNSWParams acceleratedHNSWParams; private IndexOutput hnswMeta = null, hnswVectorIndex = null; private boolean finished; @@ -86,6 +87,7 @@ public LuceneAcceleratedHNSWBinaryQuantizedVectorsWriter( this.acceleratedHNSWParams = acceleratedHNSWParams; this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; + this.graphProcessingTrace = GraphProcessingTrace.toInfoStream(infoStream, COMPONENT); vemFileName = IndexFileNames.segmentFileName( @@ -184,12 +186,16 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.BINARY, - acceleratedHNSWParams.getWriterThreads()); + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - // Write the graph to the vector index int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -284,7 +290,11 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { // Write the graph to the vector index int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java index 0291b928ad..0ad8f24c5e 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/LuceneAcceleratedHNSWScalarQuantizedVectorsWriter.java @@ -64,6 +64,7 @@ public class LuceneAcceleratedHNSWScalarQuantizedVectorsWriter extends KnnVector private final FlatVectorsWriter flatVectorsWriter; private final List fields = new ArrayList<>(); private final InfoStream infoStream; + private final GraphProcessingTrace graphProcessingTrace; private final AcceleratedHNSWParams acceleratedHNSWParams; private IndexOutput hnswMeta = null, hnswVectorIndex = null; private boolean finished; @@ -96,6 +97,7 @@ public LuceneAcceleratedHNSWScalarQuantizedVectorsWriter( this.acceleratedHNSWParams = acceleratedHNSWParams; this.flatVectorsWriter = flatVectorsWriter; this.infoStream = state.infoStream; + this.graphProcessingTrace = GraphProcessingTrace.toInfoStream(infoStream, COMPONENT); vemFileName = IndexFileNames.segmentFileName( @@ -211,12 +213,16 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw acceleratedHNSWParams.getHnswLayers(), params, QuantizationType.SCALAR, - acceleratedHNSWParams.getWriterThreads()); + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - // Write the graph to the vector index int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -310,7 +316,11 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index int[][] graphLevelNodeOffsets = - writeGraph(hnswGraph, hnswVectorIndex, acceleratedHNSWParams.getWriterThreads()); + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; // Write metadata diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/IntGraphTestMatrix.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/IntGraphTestMatrix.java new file mode 100644 index 0000000000..73ef7c5a78 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/IntGraphTestMatrix.java @@ -0,0 +1,238 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix; +import com.nvidia.cuvs.CuVSResources; +import com.nvidia.cuvs.RowView; +import java.util.Random; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +/** Minimal in-memory INT matrix used by graph materialization and serialization tests. */ +class IntGraphTestMatrix implements CuVSMatrix { + private final int[][] rows; + + IntGraphTestMatrix(int[][] rows) { + this.rows = rows; + } + + static IntGraphTestMatrix random(int rowCount, int columnCount, long seed) { + return new IntGraphTestMatrix(randomRows(rowCount, columnCount, seed)); + } + + static int[][] randomRows(int rowCount, int columnCount, long seed) { + Random random = new Random(seed); + int[][] rows = new int[rowCount][columnCount]; + for (int[] row : rows) { + for (int column = 0; column < row.length; column++) { + row[column] = random.nextInt(rowCount); + } + } + return rows; + } + + @Override + public long size() { + return rows.length; + } + + @Override + public long columns() { + return rows.length == 0 ? 0 : rows[0].length; + } + + @Override + public DataType dataType() { + return DataType.INT; + } + + @Override + public RowView getRow(long row) { + return new IntRow(rows[Math.toIntExact(row)]); + } + + @Override + public void toArray(int[][] target) { + for (int row = 0; row < rows.length; row++) { + System.arraycopy(rows[row], 0, target[row], 0, rows[row].length); + } + } + + @Override + public void toArray(float[][] target) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(byte[][] target) { + throw new UnsupportedOperationException(); + } + + @Override + public void toHost(CuVSHostMatrix target) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSHostMatrix toHost() { + throw new UnsupportedOperationException(); + } + + @Override + public void toDevice(CuVSDeviceMatrix target, CuVSResources resources) { + throw new UnsupportedOperationException(); + } + + @Override + public CuVSDeviceMatrix toDevice(CuVSResources resources) { + throw new UnsupportedOperationException(); + } + + @Override + public void close() {} + + /** Reports a synthetic device shape and fails if an unapproved host copy is attempted. */ + static class DeviceMatrix extends IntGraphTestMatrix implements CuVSDeviceMatrix { + private final long reportedColumns; + + DeviceMatrix(int[][] rows, long reportedColumns) { + super(rows); + this.reportedColumns = reportedColumns; + } + + @Override + public long columns() { + return reportedColumns; + } + + @Override + public void toHost(CuVSHostMatrix target) { + throw new AssertionError("rejected device adjacency must not be copied to host"); + } + + @Override + public CuVSHostMatrix toHost() { + throw new AssertionError("rejected device adjacency must not be copied to host"); + } + } + + static final class TrackingHostMatrix extends IntGraphTestMatrix implements CuVSHostMatrix { + private final AtomicInteger closeCount; + private final RuntimeException closeFailure; + private final ParallelExecutionProbe executionProbe; + + TrackingHostMatrix(AtomicInteger closeCount, RuntimeException closeFailure) { + this(new int[][] {{0}}, closeCount, closeFailure, null); + } + + TrackingHostMatrix( + int[][] rows, + AtomicInteger closeCount, + RuntimeException closeFailure, + ParallelExecutionProbe executionProbe) { + super(rows); + this.closeCount = closeCount; + this.closeFailure = closeFailure; + this.executionProbe = executionProbe; + } + + @Override + public int get(int row, int column) { + return getRow(row).getAsInt(column); + } + + @Override + public RowView getRow(long row) { + if (executionProbe != null) { + executionProbe.recordExecution(); + } + return super.getRow(row); + } + + @Override + public void close() { + closeCount.incrementAndGet(); + if (closeFailure != null) { + throw closeFailure; + } + } + } + + /** Holds the first operation until a second execution context reaches the same work. */ + static final class ParallelExecutionProbe { + private static final long TIMEOUT_SECONDS = 10; + + private final Set threads = ConcurrentHashMap.newKeySet(); + private final CountDownLatch parallelExecution = new CountDownLatch(1); + + void recordExecution() { + threads.add(Thread.currentThread()); + if (threads.size() > 1) { + parallelExecution.countDown(); + } + try { + if (!parallelExecution.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)) { + throw new AssertionError("work never reached a second execution context"); + } + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new AssertionError("interrupted while observing parallel work", interrupted); + } + } + + int threadCount() { + return threads.size(); + } + } + + private static final class IntRow implements RowView { + private final int[] values; + + private IntRow(int[] values) { + this.values = values; + } + + @Override + public long size() { + return values.length; + } + + @Override + public int getAsInt(long index) { + return values[Math.toIntExact(index)]; + } + + @Override + public float getAsFloat(long index) { + throw new UnsupportedOperationException(); + } + + @Override + public byte getAsByte(long index) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(int[] target) { + System.arraycopy(values, 0, target, 0, values.length); + } + + @Override + public void toArray(float[] target) { + throw new UnsupportedOperationException(); + } + + @Override + public void toArray(byte[] target) { + throw new UnsupportedOperationException(); + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java index e5818c1ba2..6794ad1f89 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWParams.java @@ -9,6 +9,7 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_CAGRA_GRAPH_BUILD_ALGO; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_CUVS_DISTANCE_TYPE; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_GRAPH_DEGREE; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_GRAPH_THREADS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_HNSW_LAYERS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_INT_GRAPH_DEGREE; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_MAX_CONN; @@ -18,6 +19,7 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.DEFAULT_WRITER_THREADS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_BEAM_WIDTH; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_GRAPH_THREADS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_HNSW_LAYERS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_INT_GRAPH_DEG; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_MAX_CONN; @@ -26,6 +28,7 @@ import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MAX_WRITER_THREADS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_BEAM_WIDTH; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_GRAPH_DEG; +import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_GRAPH_THREADS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_HNSW_LAYERS; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_INT_GRAPH_DEG; import static com.nvidia.cuvs.lucene.AcceleratedHNSWParams.MIN_MAX_CONN; @@ -55,6 +58,7 @@ public void testAcceleratedHNSWParamsDefaultValues() { AcceleratedHNSWParams params = new AcceleratedHNSWParams.Builder().build(); assertEquals(DEFAULT_BEAM_WIDTH, params.getBeamWidth()); assertEquals(DEFAULT_GRAPH_DEGREE, params.getGraphdegree()); + assertEquals(DEFAULT_GRAPH_THREADS, params.getGraphThreads()); assertEquals(DEFAULT_HNSW_LAYERS, params.getHnswLayers()); assertEquals(DEFAULT_INT_GRAPH_DEGREE, params.getIntermediateGraphDegree()); assertEquals(DEFAULT_MAX_CONN, params.getMaxConn()); @@ -208,6 +212,28 @@ public void testAcceleratedHNSWParamsInvalidWriterThreads() { } } + @Test + public void testAcceleratedHNSWParamsInvalidGraphThreads() { + for (int v : + new int[] { + random.nextInt(MIN_VALUE, MIN_GRAPH_THREADS), + random.nextInt(MAX_GRAPH_THREADS + 1, Integer.MAX_VALUE) + }) { + assertThrows( + IllegalArgumentException.class, + () -> new AcceleratedHNSWParams.Builder().withGraphThreads(v).build()); + } + } + + @Test + public void testWriterAndGraphThreadsAreIndependent() { + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder().withWriterThreads(7).withGraphThreads(3).build(); + + assertEquals(7, params.getWriterThreads()); + assertEquals(3, params.getGraphThreads()); + } + @Test public void testAcceleratedHNSWParamsInvalidNumMergeWorkers() { for (int v : diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java index f559ea3808..c4319e6f46 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCagraIndexParamsFactory.java @@ -184,6 +184,7 @@ public void testHnswHeuristicDelegatesToCuVS() { .withMaxConn(16) .withBeamWidth(100) .withWriterThreads(7) + .withGraphThreads(3) .build(); CagraIndexParams cagraParams = CagraIndexParamsFactory.create(params, 10_000, 128); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java new file mode 100644 index 0000000000..6d0d309d3f --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphCopyMemoryBudget.java @@ -0,0 +1,152 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.RamUsageEstimator; +import org.apache.lucene.util.hnsw.NeighborArray; +import org.junit.Test; + +/** Behavioral specifications for native graph-copy admission control. */ +public class TestGraphCopyMemoryBudget extends LuceneTestCase { + private static final long TIMEOUT_SECONDS = 10; + + @Test + public void reservationsFollowEstimatedPeakAcrossSupportedDegrees() { + int rows = 100; + for (int degree : new int[] {1, 32, 512}) { + long required = GraphCopyMemoryBudget.requiredHeadroom(rows, degree); + long adjacencyBytes = (long) rows * degree * Integer.BYTES; + long expected = + 2 * adjacencyBytes + + RamUsageEstimator.shallowSizeOf(new NeighborArray[rows]) + + rows + * (RamUsageEstimator.shallowSizeOfInstance(NeighborArray.class) + + RamUsageEstimator.sizeOf(new int[degree]) + + RamUsageEstimator.sizeOf(new float[degree])); + assertEquals(expected, required); + assertTrue("object layout must be included", required > 4 * adjacencyBytes); + + GraphCopyMemoryBudget exactBudget = budgetWith(required, required); + try (GraphCopyMemoryBudget.Reservation ignored = reserve(exactBudget, rows, degree)) { + assertTrue(exactBudget.tryReserve(1, 1).isEmpty()); + } + + GraphCopyMemoryBudget insufficientBudget = budgetWith(required, required - 1); + assertTrue(insufficientBudget.tryReserve(rows, degree).isEmpty()); + } + } + + @Test + public void concurrentReservationsCannotExceedSharedHeadroom() throws Exception { + long rows = 100; + long degree = 16; + long reservationBytes = GraphCopyMemoryBudget.requiredHeadroom(rows, degree); + GraphCopyMemoryBudget budget = budgetWith(2 * reservationBytes, 2 * reservationBytes); + int callers = 8; + ExecutorService executor = Executors.newFixedThreadPool(callers); + CountDownLatch start = new CountDownLatch(1); + CountDownLatch attempted = new CountDownLatch(callers); + CountDownLatch release = new CountDownLatch(1); + AtomicInteger granted = new AtomicInteger(); + List> futures = new ArrayList<>(); + try { + for (int i = 0; i < callers; i++) { + futures.add( + executor.submit( + () -> { + assertTrue(start.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + Optional reservation = + budget.tryReserve(rows, degree); + reservation.ifPresent(ignored -> granted.incrementAndGet()); + attempted.countDown(); + if (reservation.isPresent()) { + assertTrue(release.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + reservation.orElseThrow().close(); + } + return null; + })); + } + + start.countDown(); + assertTrue(attempted.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + assertEquals(2, granted.get()); + release.countDown(); + for (Future future : futures) { + future.get(TIMEOUT_SECONDS, TimeUnit.SECONDS); + } + } finally { + release.countDown(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } + } + + @Test + public void reservationIsReleasedOnFailureAndCloseIsIdempotent() { + long required = GraphCopyMemoryBudget.requiredHeadroom(100, 16); + GraphCopyMemoryBudget budget = budgetWith(required, required); + GraphCopyMemoryBudget.Reservation failedOperation = reserve(budget, 100, 16); + RuntimeException expected = new RuntimeException("expected"); + + RuntimeException actual = + assertThrows( + RuntimeException.class, + () -> { + try (failedOperation) { + throw expected; + } + }); + assertSame(expected, actual); + failedOperation.close(); + + try (GraphCopyMemoryBudget.Reservation replacement = reserve(budget, 100, 16)) { + assertTrue(budget.tryReserve(1, 1).isEmpty()); + } + } + + @Test + public void invalidOrUnavailableMemoryInformationFailsClosed() { + assertRejected(() -> null); + assertRejected(() -> new GraphCopyMemoryBudget.MemorySnapshot(0, 0)); + assertRejected(() -> new GraphCopyMemoryBudget.MemorySnapshot(1_000, -1)); + assertRejected(() -> new GraphCopyMemoryBudget.MemorySnapshot(1_000, 1_001)); + assertRejected( + () -> { + throw new UnsupportedOperationException("unavailable"); + }); + + GraphCopyMemoryBudget budget = + budgetWith(/* totalBytes= */ Long.MAX_VALUE, /* freeBytes= */ Long.MAX_VALUE); + assertTrue(budget.tryReserve(Integer.MAX_VALUE, Integer.MAX_VALUE).isEmpty()); + assertTrue(budget.tryReserve(0, 1).isEmpty()); + assertTrue(budget.tryReserve(1, 0).isEmpty()); + assertTrue(budget.tryReserve(-1, 1).isEmpty()); + assertTrue(budget.tryReserve(1, -1).isEmpty()); + } + + private static GraphCopyMemoryBudget budgetWith(long totalBytes, long freeBytes) { + return new GraphCopyMemoryBudget( + () -> new GraphCopyMemoryBudget.MemorySnapshot(totalBytes, freeBytes)); + } + + private static GraphCopyMemoryBudget.Reservation reserve( + GraphCopyMemoryBudget budget, long rows, long columns) { + return budget.tryReserve(rows, columns).orElseThrow(); + } + + private static void assertRejected(GraphCopyMemoryBudget.MemoryProbe probe) { + assertTrue(new GraphCopyMemoryBudget(probe).tryReserve(1, 1).isEmpty()); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphThreadsPersistedIndex.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphThreadsPersistedIndex.java new file mode 100644 index 0000000000..e951accc7e --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphThreadsPersistedIndex.java @@ -0,0 +1,232 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import com.carrotsearch.randomizedtesting.annotations.Name; +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Random; +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsReader; +import org.apache.lucene.codecs.hnsw.HnswGraphProvider; +import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FloatVectorValues; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReader; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.apache.lucene.tests.util.TestUtil; +import org.apache.lucene.util.InfoStream; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.junit.Test; + +@SuppressSysoutChecks(bugUrl = "") +public class TestGraphThreadsPersistedIndex extends LuceneTestCase { + + private static final String VECTOR_FIELD = "vector"; + private static final String REQUIRE_GPU_ENV = "CUVS_TESTS_REQUIRE_GPU"; + private static final int VECTOR_COUNT = GPUBuiltHnswGraph.PARALLEL_MIN_NODES + 1; + private static final int GRAPH_THREADS = 4; + + private enum WriterVariant { + FLOAT, + BINARY_QUANTIZED, + SCALAR_QUANTIZED + } + + private final WriterVariant writerVariant; + private final int dimensions; + + public TestGraphThreadsPersistedIndex( + @Name("writer") WriterVariant writerVariant, @Name("dimensions") int dimensions) { + this.writerVariant = writerVariant; + this.dimensions = dimensions; + } + + @ParametersFactory + public static List parameters() { + return List.of( + new Object[] {WriterVariant.FLOAT, 32}, + new Object[] {WriterVariant.BINARY_QUANTIZED, 129}, + new Object[] {WriterVariant.SCALAR_QUANTIZED, 32}); + } + + @Test + public void testEveryWriterForwardsGraphThreadsAndPersistsValidIndex() throws Exception { + requireCuvsSupport(); + AcceleratedHNSWParams params = + new AcceleratedHNSWParams.Builder() + .withWriterThreads(1) + .withGraphThreads(GRAPH_THREADS) + .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) + .withIntermediateGraphDegree(32) + .withGraphDegree(16) + .withHNSWLayer(1) + .build(); + Codec codec = codecFor(params); + RecordingInfoStream infoStream = new RecordingInfoStream(); + + try (Directory directory = newDirectory()) { + writeIndexWithGraphThreads(directory, codec, infoStream); + assertParallelStageObserved(infoStream, "materialization", "device-host-copy"); + assertParallelStageObserved(infoStream, "serialization", "above-threshold"); + assertPersistedGraphIsValidAndSearchable(directory); + } + } + + private void writeIndexWithGraphThreads( + Directory directory, Codec codec, RecordingInfoStream infoStream) throws Exception { + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(codec) + .setInfoStream(infoStream) + .setUseCompoundFile(false) + .setMaxBufferedDocs(VECTOR_COUNT + 1) + .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); + Random random = new Random(0x2594L); + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (int id = 0; id < VECTOR_COUNT; id++) { + float[] vector = new float[dimensions]; + for (int dimension = 0; dimension < dimensions; dimension++) { + vector[dimension] = random.nextFloat(); + } + Document document = new Document(); + document.add(new StringField("id", Integer.toString(id), Field.Store.YES)); + document.add(new KnnFloatVectorField(VECTOR_FIELD, vector, EUCLIDEAN)); + writer.addDocument(document); + } + } + } + + private static void assertPersistedGraphIsValidAndSearchable(Directory directory) + throws Exception { + TestUtil.checkIndex(directory); + try (DirectoryReader reader = DirectoryReader.open(directory)) { + assertEquals(1, reader.leaves().size()); + assertEquals(VECTOR_COUNT, reader.numDocs()); + LeafReader leaf = getOnlyLeafReader(reader); + HnswGraph graph = graphOf(leaf); + assertEquals(VECTOR_COUNT, graph.size()); + int arcs = 0; + HnswGraph.NodesIterator nodes = graph.getNodesOnLevel(0); + while (nodes.hasNext()) { + int node = nodes.nextInt(); + graph.seek(0, node); + for (int neighbor = graph.nextNeighbor(); + neighbor != NO_MORE_DOCS; + neighbor = graph.nextNeighbor()) { + assertTrue(neighbor >= 0); + assertTrue(neighbor < VECTOR_COUNT); + arcs++; + } + } + assertTrue("persisted graph contains no arcs", arcs > 0); + + int queryNode = graph.entryNode(); + assertTrue(queryNode >= 0); + assertTrue(queryNode < VECTOR_COUNT); + FloatVectorValues values = leaf.getFloatVectorValues(VECTOR_FIELD); + assertNotNull(values); + float[] query = values.vectorValue(queryNode).clone(); + int queryDoc = values.ordToDoc(queryNode); + String queryId = leaf.storedFields().document(queryDoc).get("id"); + + IndexSearcher searcher = new IndexSearcher(reader); + var hits = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, query, 10), 10); + assertEquals(10, hits.scoreDocs.length); + boolean foundQueryNode = false; + for (var hit : hits.scoreDocs) { + foundQueryNode |= queryId.equals(searcher.storedFields().document(hit.doc).get("id")); + } + assertTrue("the entry-node vector must be returned for its own query", foundQueryNode); + } + } + + private static void requireCuvsSupport() { + boolean supported = isSupported(); + if ("1".equals(System.getenv(REQUIRE_GPU_ENV))) { + assertTrue( + REQUIRE_GPU_ENV + + "=1, but cuVS failed to initialize; verify GPU visibility, matching " + + "libcuvs/libcuvs_c libraries, LD_LIBRARY_PATH, and Java native-access " + + "configuration.", + supported); + } else { + assumeTrue("cuVS not supported", supported); + } + } + + private Codec codecFor(AcceleratedHNSWParams params) throws Exception { + return switch (writerVariant) { + case FLOAT -> new Lucene101AcceleratedHNSWCodec(params); + case BINARY_QUANTIZED -> new LuceneAcceleratedHNSWBinaryQuantizedCodec(params); + case SCALAR_QUANTIZED -> new LuceneAcceleratedHNSWScalarQuantizedCodec(params); + }; + } + + private void assertParallelStageObserved( + RecordingInfoStream infoStream, String stage, String expectedReason) { + String expectedPath = + "graph-processing stage=" + + stage + + " mode=parallel reason=" + + expectedReason + + " requestedThreads=" + + GRAPH_THREADS + + " nodes=" + + VECTOR_COUNT; + assertTrue( + writerVariant + " did not report the expected path; messages: " + infoStream.messages(), + infoStream.messages().stream().anyMatch(message -> message.contains(expectedPath))); + } + + private static HnswGraph graphOf(LeafReader leaf) throws Exception { + KnnVectorsReader reader = ((CodecReader) leaf).getVectorReader(); + if (reader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { + reader = fieldsReader.getFieldReader(VECTOR_FIELD); + } + return ((HnswGraphProvider) reader).getGraph(VECTOR_FIELD); + } + + private static final class RecordingInfoStream extends InfoStream { + + private final List messages = Collections.synchronizedList(new ArrayList<>()); + + @Override + public void message(String component, String message) { + messages.add(component + ": " + message); + } + + @Override + public boolean isEnabled(String component) { + return true; + } + + @Override + public void close() {} + + List messages() { + synchronized (messages) { + return List.copyOf(messages); + } + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphWorkExecutor.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphWorkExecutor.java new file mode 100644 index 0000000000..9f8cb1c75d --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestGraphWorkExecutor.java @@ -0,0 +1,455 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicIntegerArray; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.ThreadInterruptedException; +import org.junit.Test; + +/** Behavioral specifications for the shared accelerated-HNSW graph-work scheduler. */ +public class TestGraphWorkExecutor extends LuceneTestCase { + private static final long TIMEOUT_SECONDS = 10; + + @Test + public void defaultWorkerCapReservesTheCallingThread() { + assertEquals(1, GraphWorkExecutor.defaultMaxWorkers(1)); + assertEquals(1, GraphWorkExecutor.defaultMaxWorkers(2)); + assertEquals(7, GraphWorkExecutor.defaultMaxWorkers(8)); + } + + @Test + public void saturationRunsWorkOnTheCallerWithoutLosingIt() throws Exception { + ThreadPoolExecutor executor = newExecutor(1); + CountDownLatch helperStarted = new CountDownLatch(1); + CountDownLatch releaseHelper = new CountDownLatch(1); + try { + executor.execute( + () -> { + helperStarted.countDown(); + awaitUninterruptibly(releaseHelper); + }); + assertTrue(helperStarted.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + + AtomicIntegerArray calls = new AtomicIntegerArray(32); + Set taskThreads = ConcurrentHashMap.newKeySet(); + GraphWorkExecutor.invokeAll(executor, countedTasks(calls, taskThreads)); + + assertEquals(Set.of(Thread.currentThread()), taskThreads); + assertEachCalledOnce(calls); + } finally { + releaseHelper.countDown(); + shutdown(executor); + } + } + + @Test + public void concurrentCallersUseTheProductionSharedWorkerBound() throws Exception { + int workerLimit = + GraphWorkExecutor.defaultMaxWorkers(Runtime.getRuntime().availableProcessors()); + int callers = Math.min(4, workerLimit + 1); + int tasksPerCaller = Math.ceilDiv(workerLimit + 1, callers) + 1; + ConcurrentRun run = runConcurrentWork(callers, tasksPerCaller, workerLimit); + + assertEachCalledOnce(run.calls); + assertTrue(run.maxActiveHelpers > 0); + assertTrue(run.maxActiveHelpers <= workerLimit); + } + + @Test + public void failuresWaitForStartedWorkPreserveCausesAndSkipPendingWork() throws Exception { + ThreadPoolExecutor graphExecutor = newExecutor(1); + ExecutorService callerExecutor = Executors.newSingleThreadExecutor(); + CountDownLatch helperStarted = new CountDownLatch(1); + CountDownLatch releaseHelper = new CountDownLatch(1); + CountDownLatch callerTaskFailed = new CountDownLatch(1); + AtomicInteger pendingCalls = new AtomicInteger(); + IOException helperFailure = new IOException("helper failed"); + IllegalStateException callerFailure = new IllegalStateException("caller failed"); + try { + Future invocation = + invokeAsync( + callerExecutor, + () -> + GraphWorkExecutor.invokeAll( + graphExecutor, + List.of( + () -> { + helperStarted.countDown(); + awaitUninterruptibly(releaseHelper); + throw helperFailure; + }, + () -> { + assertTrue(helperStarted.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + callerTaskFailed.countDown(); + throw callerFailure; + }, + () -> { + pendingCalls.incrementAndGet(); + return null; + }))); + + assertTrue(callerTaskFailed.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + assertFalse(invocation.isDone()); + releaseHelper.countDown(); + + Throwable thrown = invocation.get(TIMEOUT_SECONDS, TimeUnit.SECONDS); + assertTrue(thrown == helperFailure || thrown == callerFailure); + Throwable otherFailure = thrown == helperFailure ? callerFailure : helperFailure; + assertArrayEquals(new Throwable[] {otherFailure}, thrown.getSuppressed()); + assertEquals(0, pendingCalls.get()); + } finally { + releaseHelper.countDown(); + shutdown(callerExecutor); + shutdown(graphExecutor); + } + } + + @Test + public void interruptedCallerWaitsForStartedWorkAndRestoresInterrupt() throws Exception { + ThreadPoolExecutor executor = newExecutor(1); + BlockingTask helper = new BlockingTask(); + CountDownLatch callerShareFinished = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + AtomicBoolean interruptedOnExit = new AtomicBoolean(); + try { + Thread caller = + new Thread( + () -> { + try { + GraphWorkExecutor.invokeAll( + executor, + List.of( + helper, + () -> { + helper.awaitStarted(); + callerShareFinished.countDown(); + return null; + })); + } catch (Throwable thrown) { + failure.set(thrown); + } finally { + interruptedOnExit.set(Thread.currentThread().isInterrupted()); + } + }, + "graph-work-interrupted-caller"); + caller.start(); + + assertTrue(callerShareFinished.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + awaitThreadWaiting(caller); + caller.interrupt(); + awaitInterruptConsumed(caller); + assertTrue("caller returned while accepted work was still running", caller.isAlive()); + + helper.release(); + caller.join(TimeUnit.SECONDS.toMillis(TIMEOUT_SECONDS)); + assertFalse("caller did not finish", caller.isAlive()); + assertTrue(failure.get() instanceof ThreadInterruptedException); + assertTrue(interruptedOnExit.get()); + } finally { + helper.release(); + shutdown(executor); + } + } + + @Test + public void submissionFailureWaitsForAcceptedWork() throws Exception { + ThreadPoolExecutor delegate = newExecutor(1); + ExecutorService callerExecutor = Executors.newSingleThreadExecutor(); + BlockingTask acceptedTask = new BlockingTask(); + IllegalStateException expected = new IllegalStateException("submission failed"); + CountDownLatch submissionFailed = new CountDownLatch(1); + AtomicInteger submissions = new AtomicInteger(); + Executor failsAfterFirstSubmission = + task -> { + if (submissions.incrementAndGet() == 1) { + delegate.execute(task); + acceptedTask.awaitStartedUninterruptibly(); + } else { + submissionFailed.countDown(); + throw expected; + } + }; + try { + Future invocation = + invokeAsync( + callerExecutor, + () -> + GraphWorkExecutor.invokeAll( + failsAfterFirstSubmission, List.of(acceptedTask, () -> null, () -> null))); + + assertTrue(submissionFailed.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + assertFalse(invocation.isDone()); + acceptedTask.release(); + + assertSame(expected, invocation.get(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } finally { + acceptedTask.release(); + shutdown(callerExecutor); + shutdown(delegate); + } + } + + @Test + public void workersAreIsolatedDaemonsThatExpireWhenIdle() throws Exception { + ThreadPoolExecutor executor = GraphWorkExecutor.newExecutor(1, 25, TimeUnit.MILLISECONDS); + InheritableThreadLocal callerState = new InheritableThreadLocal<>(); + ClassLoader originalContextLoader = Thread.currentThread().getContextClassLoader(); + ClassLoader callerContextLoader = new ClassLoader(originalContextLoader) {}; + AtomicReference observation = new AtomicReference<>(); + CountDownLatch workerObserved = new CountDownLatch(1); + try { + callerState.set("caller-state"); + Thread.currentThread().setContextClassLoader(callerContextLoader); + Callable observeWorker = + () -> { + if (isGraphWorker()) { + Thread worker = Thread.currentThread(); + observation.set( + new WorkerObservation(worker, callerState.get(), worker.getContextClassLoader())); + workerObserved.countDown(); + } else { + assertTrue(workerObserved.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } + return null; + }; + + GraphWorkExecutor.invokeAll(executor, List.of(observeWorker, observeWorker)); + + WorkerObservation first = observation.get(); + assertNotNull(first); + assertTrue(first.thread.isDaemon()); + assertTrue(first.thread.getName().startsWith(GraphWorkExecutor.THREAD_NAME_PREFIX)); + assertNull(first.inheritedState); + assertNull(first.contextClassLoader); + + awaitPoolSize(executor, 0); + Thread replacement = invokeAndCaptureWorker(executor); + assertNotSame(first.thread, replacement); + } finally { + Thread.currentThread().setContextClassLoader(originalContextLoader); + callerState.remove(); + shutdown(executor); + } + } + + private static ConcurrentRun runConcurrentWork(int callers, int tasksPerCaller, int workerLimit) + throws Exception { + ExecutorService callerExecutor = Executors.newFixedThreadPool(callers); + AtomicIntegerArray calls = new AtomicIntegerArray(callers * tasksPerCaller); + AtomicInteger activeHelpers = new AtomicInteger(); + AtomicInteger maxActiveHelpers = new AtomicInteger(); + CountDownLatch start = new CountDownLatch(1); + CountDownLatch releaseHelpers = new CountDownLatch(1); + CountDownLatch sharedCapacityReached = new CountDownLatch(workerLimit); + List> invocations = new ArrayList<>(); + try { + for (int caller = 0; caller < callers; caller++) { + int invocationId = caller; + invocations.add( + callerExecutor.submit( + () -> { + assertTrue(start.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + List> tasks = new ArrayList<>(tasksPerCaller); + for (int task = 0; task < tasksPerCaller; task++) { + int taskId = invocationId * tasksPerCaller + task; + tasks.add( + () -> { + boolean helper = isGraphWorker(); + if (helper) { + int active = activeHelpers.incrementAndGet(); + maxActiveHelpers.accumulateAndGet(active, Math::max); + sharedCapacityReached.countDown(); + } + try { + calls.incrementAndGet(taskId); + if (helper) { + assertTrue(releaseHelpers.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } + } finally { + if (helper) { + activeHelpers.decrementAndGet(); + } + } + return null; + }); + } + GraphWorkExecutor.invokeAll(tasks); + return null; + })); + } + start.countDown(); + assertTrue(sharedCapacityReached.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + releaseHelpers.countDown(); + awaitAll(invocations); + return new ConcurrentRun(calls, maxActiveHelpers.get()); + } finally { + releaseHelpers.countDown(); + shutdown(callerExecutor); + } + } + + private static List> countedTasks( + AtomicIntegerArray calls, Set taskThreads) { + List> tasks = new ArrayList<>(calls.length()); + for (int task = 0; task < calls.length(); task++) { + int taskId = task; + tasks.add( + () -> { + calls.incrementAndGet(taskId); + taskThreads.add(Thread.currentThread()); + return null; + }); + } + return tasks; + } + + private static Future invokeAsync( + ExecutorService executor, ThrowingAction invocation) { + return executor.submit( + () -> { + try { + invocation.run(); + return null; + } catch (Throwable failure) { + return failure; + } + }); + } + + private static Thread invokeAndCaptureWorker(ThreadPoolExecutor executor) throws IOException { + AtomicReference worker = new AtomicReference<>(); + CountDownLatch observed = new CountDownLatch(1); + Callable capture = + () -> { + if (isGraphWorker()) { + worker.set(Thread.currentThread()); + observed.countDown(); + } else { + assertTrue(observed.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } + return null; + }; + GraphWorkExecutor.invokeAll(executor, List.of(capture, capture)); + assertNotNull(worker.get()); + return worker.get(); + } + + private static boolean isGraphWorker() { + return Thread.currentThread().getName().startsWith(GraphWorkExecutor.THREAD_NAME_PREFIX); + } + + private static ThreadPoolExecutor newExecutor(int maxWorkers) { + return GraphWorkExecutor.newExecutor(maxWorkers, 1, TimeUnit.MINUTES); + } + + private static void assertEachCalledOnce(AtomicIntegerArray calls) { + for (int task = 0; task < calls.length(); task++) { + assertEquals("task " + task, 1, calls.get(task)); + } + } + + private static void awaitAll(List> futures) throws Exception { + for (Future future : futures) { + future.get(TIMEOUT_SECONDS, TimeUnit.SECONDS); + } + } + + private static void awaitPoolSize(ThreadPoolExecutor executor, int expected) + throws InterruptedException { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TIMEOUT_SECONDS); + while (executor.getPoolSize() != expected && System.nanoTime() < deadline) { + Thread.sleep(10); + } + assertEquals("pool size did not reach " + expected, expected, executor.getPoolSize()); + } + + private static void awaitThreadWaiting(Thread thread) throws InterruptedException { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TIMEOUT_SECONDS); + while (thread.getState() != Thread.State.WAITING && System.nanoTime() < deadline) { + Thread.sleep(1); + } + assertEquals("thread did not enter a wait", Thread.State.WAITING, thread.getState()); + } + + private static void awaitInterruptConsumed(Thread thread) throws InterruptedException { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TIMEOUT_SECONDS); + while (thread.isInterrupted() && System.nanoTime() < deadline) { + Thread.sleep(1); + } + assertTrue("caller returned before its accepted work", thread.isAlive()); + assertFalse("caller did not consume its interrupt", thread.isInterrupted()); + } + + private static void awaitUninterruptibly(CountDownLatch latch) { + boolean interrupted = false; + while (true) { + try { + latch.await(); + break; + } catch (InterruptedException ignored) { + interrupted = true; + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + + private static void shutdown(ExecutorService executor) throws InterruptedException { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } + + private static final class BlockingTask implements Callable { + private final CountDownLatch started = new CountDownLatch(1); + private final CountDownLatch release = new CountDownLatch(1); + + @Override + public Void call() { + started.countDown(); + awaitUninterruptibly(release); + return null; + } + + private void awaitStarted() throws InterruptedException { + assertTrue(started.await(TIMEOUT_SECONDS, TimeUnit.SECONDS)); + } + + private void awaitStartedUninterruptibly() { + awaitUninterruptibly(started); + } + + private void release() { + release.countDown(); + } + } + + @FunctionalInterface + private interface ThrowingAction { + void run() throws Throwable; + } + + private record ConcurrentRun(AtomicIntegerArray calls, int maxActiveHelpers) {} + + private record WorkerObservation( + Thread thread, String inheritedState, ClassLoader contextClassLoader) {} +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java new file mode 100644 index 0000000000..4a6034b210 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphMaterialization.java @@ -0,0 +1,153 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; + +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSHostMatrix; +import com.nvidia.cuvs.CuVSMatrix; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.hnsw.HnswGraph; +import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; +import org.apache.lucene.util.hnsw.NeighborArray; +import org.junit.Test; + +/** Verifies serial and parallel CAGRA-adjacency materialization and host-copy ownership. */ +public class TestParallelGraphMaterialization extends LuceneTestCase { + + private static final int NUM_NODES = GPUBuiltHnswGraph.PARALLEL_MIN_NODES + 1000; + private static final int DEGREE = 12; + private static final int GRAPH_THREADS = 4; + + @Test + public void parallelMaterializationMatchesSerial() throws Exception { + int[][] rows = IntGraphTestMatrix.randomRows(NUM_NODES, DEGREE, 1); + IntGraphTestMatrix.ParallelExecutionProbe executionProbe = + new IntGraphTestMatrix.ParallelExecutionProbe(); + try (CuVSMatrix serialMatrix = new IntGraphTestMatrix(rows); + CuVSMatrix parallelMatrix = + new IntGraphTestMatrix.TrackingHostMatrix( + rows, new AtomicInteger(), null, executionProbe)) { + GPUBuiltHnswGraph serial = + new GPUBuiltHnswGraph( + NUM_NODES, /* dimensions= */ 4, Arrays.asList((int[]) null), List.of(serialMatrix)); + GPUBuiltHnswGraph parallel = newSingleLayerGraph(parallelMatrix, GRAPH_THREADS); + assertGraphsEqual(serial, parallel); + assertTrue(executionProbe.threadCount() > 1); + } + } + + @Test + public void overflowingDeviceShapeUsesSerialFallback() throws Exception { + int[][] adjacency = IntGraphTestMatrix.randomRows(NUM_NODES, 1, 0); + try (CuVSMatrix matrix = new IntGraphTestMatrix.DeviceMatrix(adjacency, Long.MAX_VALUE)) { + GPUBuiltHnswGraph graph = newSingleLayerGraph(matrix, GRAPH_THREADS); + assertEquals(NUM_NODES, graph.size()); + } + } + + @Test + public void failedDeviceCopyClosesHostAllocationAndSuppressesCloseFailure() { + RuntimeException copyFailure = new RuntimeException("copy failed"); + RuntimeException closeFailure = new RuntimeException("close failed"); + AtomicInteger hostCloseCount = new AtomicInteger(); + CuVSDeviceMatrix source = + new IntGraphTestMatrix.DeviceMatrix(new int[][] {{0}}, 1) { + @Override + public void toHost(CuVSHostMatrix target) { + throw copyFailure; + } + }; + CuVSHostMatrix hostCopy = + new IntGraphTestMatrix.TrackingHostMatrix(hostCloseCount, closeFailure); + + RuntimeException thrown = + assertThrows( + RuntimeException.class, () -> GPUBuiltHnswGraph.copyToHost(source, () -> hostCopy)); + + assertSame(copyFailure, thrown); + assertEquals(1, hostCloseCount.get()); + assertArrayEquals(new Throwable[] {closeFailure}, thrown.getSuppressed()); + } + + @Test + public void admittedDeviceCopyIsMaterializedInParallelAndReleased() throws Exception { + int[][] expectedRows = IntGraphTestMatrix.randomRows(NUM_NODES, DEGREE, 3); + int[][] wrongSourceRows = new int[NUM_NODES][DEGREE]; + AtomicInteger hostCloseCount = new AtomicInteger(); + AtomicInteger copyCount = new AtomicInteger(); + IntGraphTestMatrix.ParallelExecutionProbe executionProbe = + new IntGraphTestMatrix.ParallelExecutionProbe(); + CuVSHostMatrix hostCopy = + new IntGraphTestMatrix.TrackingHostMatrix( + expectedRows, hostCloseCount, null, executionProbe); + CuVSDeviceMatrix source = + new IntGraphTestMatrix.DeviceMatrix(wrongSourceRows, DEGREE) { + @Override + public void toHost(CuVSHostMatrix target) { + assertSame(hostCopy, target); + copyCount.incrementAndGet(); + } + }; + long requiredHeadroom = GraphCopyMemoryBudget.requiredHeadroom(NUM_NODES, DEGREE); + GraphCopyMemoryBudget budget = + new GraphCopyMemoryBudget( + () -> new GraphCopyMemoryBudget.MemorySnapshot(requiredHeadroom, requiredHeadroom)); + + NeighborArray[] neighbors = + GPUBuiltHnswGraph.materializeDeviceAdjacency( + source, NUM_NODES, GRAPH_THREADS, budget, () -> hostCopy); + + for (int node = 0; node < NUM_NODES; node++) { + assertArrayEquals( + expectedRows[node], Arrays.copyOf(neighbors[node].nodes(), neighbors[node].size())); + } + assertEquals(1, copyCount.get()); + assertEquals(1, hostCloseCount.get()); + assertTrue(executionProbe.threadCount() > 1); + try (GraphCopyMemoryBudget.Reservation ignored = + budget.tryReserve(NUM_NODES, DEGREE).orElseThrow()) { + // The first reservation was released after materialization. + } + } + + private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, int numThreads) + throws IOException { + return new GPUBuiltHnswGraph( + NUM_NODES, + /* dimensions= */ 4, + Arrays.asList((int[]) null), + List.of(layer0Adjacency), + numThreads); + } + + private static void assertGraphsEqual(HnswGraph a, HnswGraph b) throws Exception { + assertEquals(a.numLevels(), b.numLevels()); + for (int level = 0; level < a.numLevels(); level++) { + int[] nodes = NodesIterator.getSortedNodes(a.getNodesOnLevel(level)); + for (int node : nodes) { + assertArrayEquals( + "node " + node + " at level " + level + " has different neighbors", + arcsOf(a, level, node), + arcsOf(b, level, node)); + } + } + } + + private static int[] arcsOf(HnswGraph graph, int level, int node) throws Exception { + graph.seek(level, node); + List arcs = new ArrayList<>(); + for (int n = graph.nextNeighbor(); n != NO_MORE_DOCS; n = graph.nextNeighbor()) { + arcs.add(n); + } + return arcs.stream().mapToInt(Integer::intValue).toArray(); + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java new file mode 100644 index 0000000000..ac6c1fecc4 --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestParallelGraphSerialization.java @@ -0,0 +1,259 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import com.nvidia.cuvs.CuVSMatrix; +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.IOContext; +import org.apache.lucene.store.IndexInput; +import org.apache.lucene.store.IndexOutput; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.util.hnsw.NeighborArray; +import org.junit.Test; + +/** Verifies parallel level-zero graph serialization is byte-identical to serial serialization. */ +public class TestParallelGraphSerialization extends LuceneTestCase { + + private static final int NUM_NODES = GPUBuiltHnswGraph.PARALLEL_MIN_NODES + 1000; + private static final int DEGREE = 12; + private static final int GRAPH_THREADS = 4; + + @Test + public void parallelSerializationMatchesSerial() throws Exception { + try (CuVSMatrix matrix = IntGraphTestMatrix.random(NUM_NODES, DEGREE, 2); + Directory dir = new ByteBuffersDirectory()) { + GPUBuiltHnswGraph serialGraph = newSingleLayerGraph(matrix); + IntGraphTestMatrix.ParallelExecutionProbe executionProbe = + new IntGraphTestMatrix.ParallelExecutionProbe(); + GPUBuiltHnswGraph parallelGraph = new RecordingGraph(matrix, executionProbe); + assertSerialAndParallelMatch(serialGraph, parallelGraph, dir); + assertTrue(executionProbe.threadCount() > 1); + } + } + + @Test + public void parallelSerializationMatchesSerialAcrossByteBoundedWave() throws Exception { + int maxConn = 512; + int waveNodes = AcceleratedHNSWUtils.serializationWaveNodes(maxConn); + int numNodes = Math.max(GPUBuiltHnswGraph.PARALLEL_MIN_NODES + 1, waveNodes + 2); + long worstCaseNodeBytes = 5L + maxConn * 5L; + IntGraphTestMatrix.ParallelExecutionProbe executionProbe = + new IntGraphTestMatrix.ParallelExecutionProbe(); + + try (Directory dir = new ByteBuffersDirectory()) { + assertTrue(waveNodes < AcceleratedHNSWUtils.MAX_SERIALIZATION_WAVE_NODES); + assertTrue( + waveNodes * worstCaseNodeBytes <= AcceleratedHNSWUtils.MAX_SERIALIZED_BYTES_PER_WAVE); + assertTrue( + (waveNodes + 1L) * worstCaseNodeBytes + > AcceleratedHNSWUtils.MAX_SERIALIZED_BYTES_PER_WAVE); + assertSerialAndParallelMatch( + new LazyBoundaryGraph(numNodes, maxConn, waveNodes, null), + new LazyBoundaryGraph(numNodes, maxConn, waveNodes, executionProbe), + dir); + assertTrue(executionProbe.threadCount() > 1); + } + } + + @Test + public void lowDegreeSerializationWaveUsesAbsoluteNodeCap() { + assertEquals( + AcceleratedHNSWUtils.MAX_SERIALIZATION_WAVE_NODES, + AcceleratedHNSWUtils.serializationWaveNodes(/* maxConn= */ 0)); + } + + @Test + public void serialSerializationRejectsMissingAdjacency() throws Exception { + assertMissingAdjacencyRejected(/* graphSize= */ 1, /* graphThreads= */ 1); + } + + @Test + public void parallelSerializationRejectsMissingAdjacency() throws Exception { + assertMissingAdjacencyRejected( + GPUBuiltHnswGraph.PARALLEL_MIN_NODES, /* graphThreads= */ GRAPH_THREADS); + } + + private static void assertMissingAdjacencyRejected(int graphSize, int graphThreads) + throws Exception { + try (Directory dir = new ByteBuffersDirectory(); + IndexOutput out = dir.createOutput("missing-adjacency", IOContext.DEFAULT)) { + GPUBuiltHnswGraph graph = new MissingAdjacencyGraph(graphSize); + expectThrows( + NullPointerException.class, + () -> AcceleratedHNSWUtils.writeGraph(graph, out, graphThreads)); + } + } + + private static void assertSerialAndParallelMatch( + GPUBuiltHnswGraph serialGraph, GPUBuiltHnswGraph parallelGraph, Directory dir) + throws Exception { + int[][] serialOffsets; + try (IndexOutput out = dir.createOutput("serial", IOContext.DEFAULT)) { + serialOffsets = AcceleratedHNSWUtils.writeGraph(serialGraph, out); + } + int[][] parallelOffsets; + try (IndexOutput out = dir.createOutput("parallel", IOContext.DEFAULT)) { + parallelOffsets = AcceleratedHNSWUtils.writeGraph(parallelGraph, out, GRAPH_THREADS); + } + + assertEquals(serialOffsets.length, parallelOffsets.length); + for (int level = 0; level < serialOffsets.length; level++) { + assertArrayEquals(serialOffsets[level], parallelOffsets[level]); + } + assertArrayEquals(readAllBytes(dir, "serial"), readAllBytes(dir, "parallel")); + } + + private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency) { + return new GPUBuiltHnswGraph( + NUM_NODES, /* dimensions= */ 4, Arrays.asList((int[]) null), List.of(layer0Adjacency)); + } + + private static byte[] readAllBytes(Directory dir, String name) throws Exception { + try (IndexInput in = dir.openInput(name, IOContext.DEFAULT)) { + byte[] bytes = new byte[(int) in.length()]; + in.readBytes(bytes, 0, bytes.length); + return bytes; + } + } + + private static final class RecordingGraph extends GPUBuiltHnswGraph { + private final IntGraphTestMatrix.ParallelExecutionProbe executionProbe; + + RecordingGraph( + CuVSMatrix layer0Adjacency, IntGraphTestMatrix.ParallelExecutionProbe executionProbe) { + super(NUM_NODES, /* dimensions= */ 4, Arrays.asList((int[]) null), List.of(layer0Adjacency)); + this.executionProbe = executionProbe; + } + + @Override + public NeighborArray getNeighbors(int level, int node) { + executionProbe.recordExecution(); + return super.getNeighbors(level, node); + } + } + + /** + * Crosses a maxConn-derived byte-bounded wave with sparse rows and no retained heap graph. + */ + private static final class LazyBoundaryGraph extends GPUBuiltHnswGraph { + private static final NeighborArray EMPTY_NEIGHBORS = new NeighborArray(0, true); + + private final int graphSize; + private final int maxConn; + private final int waveNodes; + private final IntGraphTestMatrix.ParallelExecutionProbe executionProbe; + + LazyBoundaryGraph( + int graphSize, + int maxConn, + int waveNodes, + IntGraphTestMatrix.ParallelExecutionProbe executionProbe) + throws IOException { + super( + 0, + /* dimensions= */ 4, + Arrays.asList((int[]) null), + List.of(new IntGraphTestMatrix(new int[0][])), + 1); + this.graphSize = graphSize; + this.maxConn = maxConn; + this.waveNodes = waveNodes; + this.executionProbe = executionProbe; + } + + @Override + public int size() { + return graphSize; + } + + @Override + public int maxConn() { + return maxConn; + } + + @Override + public NodesIterator getNodesOnLevel(int level) { + return new RangeNodesIterator(level == 0 ? graphSize : 0); + } + + @Override + public NeighborArray getNeighbors(int level, int node) { + if (executionProbe != null) { + executionProbe.recordExecution(); + } + if (node < waveNodes - 1) { + return EMPTY_NEIGHBORS; + } + NeighborArray neighbors = new NeighborArray(1, true); + neighbors.addInOrder(node, 1.0f); + return neighbors; + } + } + + private static final class MissingAdjacencyGraph extends GPUBuiltHnswGraph { + private final int graphSize; + + MissingAdjacencyGraph(int graphSize) throws IOException { + super( + 0, + /* dimensions= */ 4, + Arrays.asList((int[]) null), + List.of(new IntGraphTestMatrix(new int[0][])), + 1); + this.graphSize = graphSize; + } + + @Override + public int size() { + return graphSize; + } + + @Override + public int maxConn() { + return 0; + } + + @Override + public NodesIterator getNodesOnLevel(int level) { + return new RangeNodesIterator(level == 0 ? graphSize : 0); + } + + @Override + public NeighborArray getNeighbors(int level, int node) { + return null; + } + } + + private static final class RangeNodesIterator extends GPUBuiltHnswGraph.NodesIterator { + private int current = -1; + + RangeNodesIterator(int size) { + super(size); + } + + @Override + public boolean hasNext() { + return current + 1 < size; + } + + @Override + public int nextInt() { + return ++current; + } + + @Override + public int consume(int[] dest) { + int count = Math.min(dest.length, size - (current + 1)); + for (int i = 0; i < count; i++) { + dest[i] = ++current; + } + return count; + } + } +} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java deleted file mode 100644 index ad3366dd14..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphMaterialization.java +++ /dev/null @@ -1,307 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; - -import com.nvidia.cuvs.CuVSDeviceMatrix; -import com.nvidia.cuvs.CuVSHostMatrix; -import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.RowView; -import java.io.IOException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; -import java.util.Random; -import java.util.concurrent.atomic.AtomicInteger; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.util.hnsw.HnswGraph; -import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; -import org.junit.Test; - -/** Verifies serial and parallel CAGRA-adjacency materialization are equivalent and bounded. */ -public class TestWriterThreadsGraphMaterialization extends LuceneTestCase { - - private static final int NUM_NODES = GPUBuiltHnswGraph.PARALLEL_MIN_NODES + 1000; - private static final int DEGREE = 12; - private static final int NUM_THREADS = 4; - - @Test - public void parallelMaterializationMatchesSerial() throws Exception { - int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(1)); - - try (CuVSMatrix matrix = new ArrayMatrix(adjacency)) { - GPUBuiltHnswGraph serial = - new GPUBuiltHnswGraph( - NUM_NODES, /* dimensions= */ 4, Arrays.asList((int[]) null), List.of(matrix)); - GPUBuiltHnswGraph parallel = newSingleLayerGraph(matrix, NUM_THREADS); - assertGraphsEqual(serial, parallel); - } - } - - @Test - public void graphCopyBudgetHandlesExpectedDatasetSizesAndOverflow() { - assertTrue(GPUBuiltHnswGraph.fitsParallelGraphCopyBudget(25_000_000L, 32)); - assertFalse(GPUBuiltHnswGraph.fitsParallelGraphCopyBudget(100_000_000L, 32)); - assertFalse(GPUBuiltHnswGraph.fitsParallelGraphCopyBudget(Long.MAX_VALUE, Long.MAX_VALUE)); - } - - @Test - public void oversizedDeviceAdjacencyUsesSerialFallback() throws Exception { - int[][] adjacency = randomAdjacency(NUM_NODES, 1, new Random(0)); - long oversizedColumns = - GPUBuiltHnswGraph.MAX_PARALLEL_GRAPH_COPY_BYTES / Integer.BYTES / NUM_NODES + 1; - try (CuVSMatrix matrix = new ArrayDeviceMatrix(adjacency, oversizedColumns)) { - GPUBuiltHnswGraph graph = newSingleLayerGraph(matrix, NUM_THREADS); - assertEquals(NUM_NODES, graph.size()); - } - } - - @Test - public void failedDeviceCopyClosesHostAllocationAndSuppressesCloseFailure() { - RuntimeException copyFailure = new RuntimeException("copy failed"); - RuntimeException closeFailure = new RuntimeException("close failed"); - AtomicInteger hostCloseCount = new AtomicInteger(); - CuVSDeviceMatrix source = - new ArrayDeviceMatrix(new int[][] {{0}}, 1) { - @Override - public void toHost(CuVSHostMatrix target) { - throw copyFailure; - } - }; - CuVSHostMatrix hostCopy = new TrackingHostMatrix(hostCloseCount, closeFailure); - - RuntimeException thrown = - assertThrows( - RuntimeException.class, () -> GPUBuiltHnswGraph.copyToHost(source, () -> hostCopy)); - - assertSame(copyFailure, thrown); - assertEquals(1, hostCloseCount.get()); - assertArrayEquals(new Throwable[] {closeFailure}, thrown.getSuppressed()); - } - - @Test - public void successfulDeviceCopyTransfersHostOwnershipToCaller() { - AtomicInteger hostCloseCount = new AtomicInteger(); - CuVSHostMatrix hostCopy = new TrackingHostMatrix(hostCloseCount, null); - CuVSDeviceMatrix source = - new ArrayDeviceMatrix(new int[][] {{0}}, 1) { - @Override - public void toHost(CuVSHostMatrix target) { - assertSame(hostCopy, target); - } - }; - - CuVSHostMatrix returned = GPUBuiltHnswGraph.copyToHost(source, () -> hostCopy); - - assertSame(hostCopy, returned); - assertEquals(0, hostCloseCount.get()); - returned.close(); - assertEquals(1, hostCloseCount.get()); - } - - private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency, int numThreads) - throws IOException { - return new GPUBuiltHnswGraph( - NUM_NODES, - /* dimensions= */ 4, - Arrays.asList((int[]) null), - List.of(layer0Adjacency), - numThreads); - } - - private static void assertGraphsEqual(HnswGraph a, HnswGraph b) throws Exception { - assertEquals(a.numLevels(), b.numLevels()); - for (int level = 0; level < a.numLevels(); level++) { - int[] nodes = NodesIterator.getSortedNodes(a.getNodesOnLevel(level)); - for (int node : nodes) { - assertArrayEquals( - "node " + node + " at level " + level + " has different neighbors", - arcsOf(a, level, node), - arcsOf(b, level, node)); - } - } - } - - private static int[] arcsOf(HnswGraph graph, int level, int node) throws Exception { - graph.seek(level, node); - List arcs = new ArrayList<>(); - for (int n = graph.nextNeighbor(); n != NO_MORE_DOCS; n = graph.nextNeighbor()) { - arcs.add(n); - } - return arcs.stream().mapToInt(Integer::intValue).toArray(); - } - - private static int[][] randomAdjacency(int numNodes, int degree, Random random) { - int[][] adjacency = new int[numNodes][degree]; - for (int[] row : adjacency) { - for (int j = 0; j < degree; j++) { - row[j] = random.nextInt(numNodes); - } - } - return adjacency; - } - - private static class ArrayMatrix implements CuVSMatrix { - private final int[][] rows; - - ArrayMatrix(int[][] rows) { - this.rows = rows; - } - - @Override - public long size() { - return rows.length; - } - - @Override - public long columns() { - return rows.length == 0 ? 0 : rows[0].length; - } - - @Override - public DataType dataType() { - return DataType.INT; - } - - @Override - public RowView getRow(long row) { - return new ArrayRow(rows[Math.toIntExact(row)]); - } - - @Override - public void toArray(int[][] target) { - for (int i = 0; i < rows.length; i++) { - System.arraycopy(rows[i], 0, target[i], 0, rows[i].length); - } - } - - @Override - public void toArray(float[][] target) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(byte[][] target) { - throw new UnsupportedOperationException(); - } - - @Override - public void toHost(CuVSHostMatrix target) { - throw new UnsupportedOperationException(); - } - - @Override - public CuVSHostMatrix toHost() { - throw new UnsupportedOperationException(); - } - - @Override - public void toDevice(CuVSDeviceMatrix target, CuVSResources resources) { - throw new UnsupportedOperationException(); - } - - @Override - public CuVSDeviceMatrix toDevice(CuVSResources resources) { - throw new UnsupportedOperationException(); - } - - @Override - public void close() {} - } - - /** Reports an oversized device shape and fails if the guarded host-copy path is reached. */ - private static class ArrayDeviceMatrix extends ArrayMatrix implements CuVSDeviceMatrix { - private final long reportedColumns; - - ArrayDeviceMatrix(int[][] rows, long reportedColumns) { - super(rows); - this.reportedColumns = reportedColumns; - } - - @Override - public long columns() { - return reportedColumns; - } - - @Override - public void toHost(CuVSHostMatrix target) { - throw new AssertionError("oversized device adjacency must not be copied to host"); - } - - @Override - public CuVSHostMatrix toHost() { - throw new AssertionError("oversized device adjacency must not be copied to host"); - } - } - - private static final class TrackingHostMatrix extends ArrayMatrix implements CuVSHostMatrix { - private final AtomicInteger closeCount; - private final RuntimeException closeFailure; - - TrackingHostMatrix(AtomicInteger closeCount, RuntimeException closeFailure) { - super(new int[][] {{0}}); - this.closeCount = closeCount; - this.closeFailure = closeFailure; - } - - @Override - public int get(int row, int column) { - return 0; - } - - @Override - public void close() { - closeCount.incrementAndGet(); - if (closeFailure != null) { - throw closeFailure; - } - } - } - - private static final class ArrayRow implements RowView { - private final int[] values; - - ArrayRow(int[] values) { - this.values = values; - } - - @Override - public long size() { - return values.length; - } - - @Override - public int getAsInt(long index) { - return values[Math.toIntExact(index)]; - } - - @Override - public float getAsFloat(long index) { - throw new UnsupportedOperationException(); - } - - @Override - public byte getAsByte(long index) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(int[] target) { - System.arraycopy(values, 0, target, 0, values.length); - } - - @Override - public void toArray(float[] target) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(byte[] target) { - throw new UnsupportedOperationException(); - } - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java deleted file mode 100644 index 2d7188b3a1..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsGraphSerialization.java +++ /dev/null @@ -1,261 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import com.nvidia.cuvs.CuVSDeviceMatrix; -import com.nvidia.cuvs.CuVSHostMatrix; -import com.nvidia.cuvs.CuVSMatrix; -import com.nvidia.cuvs.CuVSResources; -import com.nvidia.cuvs.RowView; -import java.io.IOException; -import java.util.Arrays; -import java.util.List; -import java.util.Random; -import org.apache.lucene.store.ByteBuffersDirectory; -import org.apache.lucene.store.Directory; -import org.apache.lucene.store.IOContext; -import org.apache.lucene.store.IndexInput; -import org.apache.lucene.store.IndexOutput; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.util.hnsw.NeighborArray; -import org.junit.Test; - -/** Verifies parallel level-zero graph serialization is byte-identical to serial serialization. */ -public class TestWriterThreadsGraphSerialization extends LuceneTestCase { - - private static final int NUM_NODES = AcceleratedHNSWUtils.PARALLEL_MIN_NODES + 1000; - private static final int DEGREE = 12; - private static final int NUM_THREADS = 4; - - @Test - public void parallelSerializationMatchesSerial() throws Exception { - int[][] adjacency = randomAdjacency(NUM_NODES, DEGREE, new Random(2)); - - try (CuVSMatrix matrix = new ArrayMatrix(adjacency); - Directory dir = new ByteBuffersDirectory()) { - GPUBuiltHnswGraph serialGraph = newSingleLayerGraph(matrix); - GPUBuiltHnswGraph parallelGraph = newSingleLayerGraph(matrix); - assertSerialAndParallelMatch(serialGraph, parallelGraph, dir); - } - } - - @Test - public void parallelSerializationMatchesSerialAcrossFixedWaveBoundary() throws Exception { - int numNodes = AcceleratedHNSWUtils.SERIALIZATION_WAVE_NODES + 1; - assertTrue(numNodes > AcceleratedHNSWUtils.SERIALIZATION_WAVE_NODES); - - try (Directory dir = new ByteBuffersDirectory()) { - assertSerialAndParallelMatch(new LazyEmptyGraph(numNodes), new LazyEmptyGraph(numNodes), dir); - } - } - - private static void assertSerialAndParallelMatch( - GPUBuiltHnswGraph serialGraph, GPUBuiltHnswGraph parallelGraph, Directory dir) - throws Exception { - int[][] serialOffsets; - try (IndexOutput out = dir.createOutput("serial", IOContext.DEFAULT)) { - serialOffsets = AcceleratedHNSWUtils.writeGraph(serialGraph, out); - } - int[][] parallelOffsets; - try (IndexOutput out = dir.createOutput("parallel", IOContext.DEFAULT)) { - parallelOffsets = AcceleratedHNSWUtils.writeGraph(parallelGraph, out, NUM_THREADS); - } - - assertEquals(serialOffsets.length, parallelOffsets.length); - for (int level = 0; level < serialOffsets.length; level++) { - assertArrayEquals(serialOffsets[level], parallelOffsets[level]); - } - assertArrayEquals(readAllBytes(dir, "serial"), readAllBytes(dir, "parallel")); - } - - private static GPUBuiltHnswGraph newSingleLayerGraph(CuVSMatrix layer0Adjacency) { - return new GPUBuiltHnswGraph( - NUM_NODES, /* dimensions= */ 4, Arrays.asList((int[]) null), List.of(layer0Adjacency)); - } - - private static byte[] readAllBytes(Directory dir, String name) throws Exception { - try (IndexInput in = dir.openInput(name, IOContext.DEFAULT)) { - byte[] bytes = new byte[(int) in.length()]; - in.readBytes(bytes, 0, bytes.length); - return bytes; - } - } - - private static int[][] randomAdjacency(int numNodes, int degree, Random random) { - int[][] adjacency = new int[numNodes][degree]; - for (int[] row : adjacency) { - for (int j = 0; j < degree; j++) { - row[j] = random.nextInt(numNodes); - } - } - return adjacency; - } - - /** Supplies an empty graph lazily so the fixed wave boundary can be tested with little heap. */ - private static final class LazyEmptyGraph extends GPUBuiltHnswGraph { - private final int graphSize; - - LazyEmptyGraph(int graphSize) throws IOException { - super( - 0, - /* dimensions= */ 4, - Arrays.asList((int[]) null), - List.of(new ArrayMatrix(new int[0][])), - 1); - this.graphSize = graphSize; - } - - @Override - public int size() { - return graphSize; - } - - @Override - public int maxConn() { - return 0; - } - - @Override - public NodesIterator getNodesOnLevel(int level) { - return new RangeNodesIterator(level == 0 ? graphSize : 0); - } - - @Override - public NeighborArray getNeighbors(int level, int node) { - return null; - } - } - - private static final class RangeNodesIterator extends GPUBuiltHnswGraph.NodesIterator { - private int current = -1; - - RangeNodesIterator(int size) { - super(size); - } - - @Override - public boolean hasNext() { - return current + 1 < size; - } - - @Override - public int nextInt() { - return ++current; - } - - @Override - public int consume(int[] dest) { - int count = Math.min(dest.length, size - (current + 1)); - for (int i = 0; i < count; i++) { - dest[i] = ++current; - } - return count; - } - } - - private static final class ArrayMatrix implements CuVSMatrix { - private final int[][] rows; - - ArrayMatrix(int[][] rows) { - this.rows = rows; - } - - @Override - public long size() { - return rows.length; - } - - @Override - public long columns() { - return rows.length == 0 ? 0 : rows[0].length; - } - - @Override - public DataType dataType() { - return DataType.INT; - } - - @Override - public RowView getRow(long row) { - int[] values = rows[Math.toIntExact(row)]; - return new RowView() { - @Override - public long size() { - return values.length; - } - - @Override - public int getAsInt(long index) { - return values[Math.toIntExact(index)]; - } - - @Override - public float getAsFloat(long index) { - throw new UnsupportedOperationException(); - } - - @Override - public byte getAsByte(long index) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(int[] target) { - System.arraycopy(values, 0, target, 0, values.length); - } - - @Override - public void toArray(float[] target) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(byte[] target) { - throw new UnsupportedOperationException(); - } - }; - } - - @Override - public void toArray(int[][] target) { - for (int i = 0; i < rows.length; i++) { - System.arraycopy(rows[i], 0, target[i], 0, rows[i].length); - } - } - - @Override - public void toArray(float[][] target) { - throw new UnsupportedOperationException(); - } - - @Override - public void toArray(byte[][] target) { - throw new UnsupportedOperationException(); - } - - @Override - public void toHost(CuVSHostMatrix target) { - throw new UnsupportedOperationException(); - } - - @Override - public CuVSHostMatrix toHost() { - throw new UnsupportedOperationException(); - } - - @Override - public void toDevice(CuVSDeviceMatrix target, CuVSResources resources) { - throw new UnsupportedOperationException(); - } - - @Override - public CuVSDeviceMatrix toDevice(CuVSResources resources) { - throw new UnsupportedOperationException(); - } - - @Override - public void close() {} - } -} diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsPersistedIndex.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsPersistedIndex.java deleted file mode 100644 index 6a5fcb32f1..0000000000 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestWriterThreadsPersistedIndex.java +++ /dev/null @@ -1,126 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -package com.nvidia.cuvs.lucene; - -import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; -import static org.apache.lucene.index.VectorSimilarityFunction.EUCLIDEAN; -import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; - -import java.util.Random; -import org.apache.lucene.codecs.Codec; -import org.apache.lucene.codecs.KnnVectorsReader; -import org.apache.lucene.codecs.hnsw.HnswGraphProvider; -import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; -import org.apache.lucene.document.Document; -import org.apache.lucene.document.Field; -import org.apache.lucene.document.KnnFloatVectorField; -import org.apache.lucene.document.StringField; -import org.apache.lucene.index.CodecReader; -import org.apache.lucene.index.DirectoryReader; -import org.apache.lucene.index.FloatVectorValues; -import org.apache.lucene.index.IndexWriter; -import org.apache.lucene.index.IndexWriterConfig; -import org.apache.lucene.index.LeafReader; -import org.apache.lucene.search.IndexSearcher; -import org.apache.lucene.search.KnnFloatVectorQuery; -import org.apache.lucene.store.Directory; -import org.apache.lucene.tests.util.LuceneTestCase; -import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; -import org.apache.lucene.tests.util.TestUtil; -import org.apache.lucene.util.hnsw.HnswGraph; -import org.junit.Test; - -@SuppressSysoutChecks(bugUrl = "") -public class TestWriterThreadsPersistedIndex extends LuceneTestCase { - - private static final String VECTOR_FIELD = "vector"; - private static final int VECTOR_COUNT = AcceleratedHNSWUtils.PARALLEL_MIN_NODES + 1; - private static final int DIMENSIONS = 32; - - @Test - public void testParallelGraphRoundTripAboveThreshold() throws Exception { - assumeTrue("cuVS not supported", isSupported()); - AcceleratedHNSWParams params = - new AcceleratedHNSWParams.Builder() - .withWriterThreads(4) - .withStrategy(AcceleratedHNSWParams.Strategy.CUSTOM) - .withIntermediateGraphDegree(32) - .withGraphDegree(16) - .withHNSWLayer(1) - .build(); - Codec codec = new Lucene101AcceleratedHNSWCodec(params); - Random random = new Random(0x2594L); - - try (Directory directory = newDirectory()) { - IndexWriterConfig config = - new IndexWriterConfig() - .setCodec(codec) - .setUseCompoundFile(false) - .setMaxBufferedDocs(VECTOR_COUNT + 1) - .setRAMBufferSizeMB(IndexWriterConfig.DISABLE_AUTO_FLUSH); - try (IndexWriter writer = new IndexWriter(directory, config)) { - for (int id = 0; id < VECTOR_COUNT; id++) { - float[] vector = new float[DIMENSIONS]; - for (int dimension = 0; dimension < DIMENSIONS; dimension++) { - vector[dimension] = random.nextFloat(); - } - Document document = new Document(); - document.add(new StringField("id", Integer.toString(id), Field.Store.YES)); - document.add(new KnnFloatVectorField(VECTOR_FIELD, vector, EUCLIDEAN)); - writer.addDocument(document); - } - } - - TestUtil.checkIndex(directory); - try (DirectoryReader reader = DirectoryReader.open(directory)) { - assertEquals(1, reader.leaves().size()); - assertEquals(VECTOR_COUNT, reader.numDocs()); - LeafReader leaf = getOnlyLeafReader(reader); - HnswGraph graph = graphOf(leaf); - assertEquals(VECTOR_COUNT, graph.size()); - int arcs = 0; - HnswGraph.NodesIterator nodes = graph.getNodesOnLevel(0); - while (nodes.hasNext()) { - int node = nodes.nextInt(); - graph.seek(0, node); - for (int neighbor = graph.nextNeighbor(); - neighbor != NO_MORE_DOCS; - neighbor = graph.nextNeighbor()) { - assertTrue(neighbor >= 0); - assertTrue(neighbor < VECTOR_COUNT); - arcs++; - } - } - assertTrue("persisted graph contains no arcs", arcs > 0); - - int queryNode = graph.entryNode(); - assertTrue(queryNode >= 0); - assertTrue(queryNode < VECTOR_COUNT); - FloatVectorValues values = leaf.getFloatVectorValues(VECTOR_FIELD); - assertNotNull(values); - float[] query = values.vectorValue(queryNode).clone(); - int queryDoc = values.ordToDoc(queryNode); - String queryId = leaf.storedFields().document(queryDoc).get("id"); - - IndexSearcher searcher = new IndexSearcher(reader); - var hits = searcher.search(new KnnFloatVectorQuery(VECTOR_FIELD, query, 10), 10); - assertEquals(10, hits.scoreDocs.length); - boolean foundQueryNode = false; - for (var hit : hits.scoreDocs) { - foundQueryNode |= queryId.equals(searcher.storedFields().document(hit.doc).get("id")); - } - assertTrue("the entry-node vector must be returned for its own query", foundQueryNode); - } - } - } - - private static HnswGraph graphOf(LeafReader leaf) throws Exception { - KnnVectorsReader reader = ((CodecReader) leaf).getVectorReader(); - if (reader instanceof PerFieldKnnVectorsFormat.FieldsReader fieldsReader) { - reader = fieldsReader.getFieldReader(VECTOR_FIELD); - } - return ((HnswGraphProvider) reader).getGraph(VECTOR_FIELD); - } -}