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 27ffe5cb39..692f43e77c 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-acceleratedhnswparams.md @@ -18,13 +18,28 @@ public class AcceleratedHNSWParams public int getWriterThreads() ``` -Get the cuVS writer threads parameter +Get the native cuVS writer threads parameter. **Returns** cuVS writer threads parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:149`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:157`_ + +### 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:167`_ ### getIntermediateGraphDegree @@ -38,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:158`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:176`_ ### getGraphdegree @@ -52,7 +67,7 @@ Get the graph degree the graph degree parameter -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:167`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:185`_ ### getHnswLayers @@ -66,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:176`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:194`_ ### getMaxConn @@ -80,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:185`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:203`_ ### getBeamWidth @@ -94,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:194`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:212`_ ### getCagraGraphBuildAlgo @@ -108,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:203`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:221`_ ### getCuVSIvfPqParams @@ -122,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:212`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:230`_ ### getNumMergeWorkers @@ -136,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:221`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:239`_ ### getMergeExec @@ -150,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:230`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:248`_ ### getStrategy @@ -167,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:242`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:260`_ ### getCuvsDistanceType @@ -181,7 +196,7 @@ Get the cuvs distance type the distance type -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:251`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:269`_ ### getNNDescentNumIterations @@ -195,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:260`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:278`_ ### getHnswHeuristicType @@ -210,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:270`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:288`_ ### withWriterThreads @@ -218,7 +233,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWP public Builder withWriterThreads(int writerThreads) ``` -Set the number of cuVS writer threads while building the index +Set the 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\} @@ -232,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:356`_ + +### 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:369`_ ### withIntermediateGraphDegree @@ -254,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:382`_ ### withGraphDegree @@ -276,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:395`_ ### withHNSWLayer @@ -298,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:408`_ ### withMaxConn @@ -320,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:421`_ ### withBeamWidth @@ -342,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:434`_ ### withCagraGraphBuildAlgo @@ -363,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:446`_ ### withCuVSIvfPqParams @@ -383,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:457`_ ### withNumMergeWorkers @@ -404,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:469`_ ### withMergeExecutorService @@ -425,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:481`_ ### withStrategy @@ -451,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:498`_ ### withCuvsDistanceType @@ -471,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:509`_ ### withNNDescentNumIterations @@ -494,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:523`_ ### withHnswHeuristicType @@ -518,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:538`_ ### build @@ -532,6 +569,6 @@ Create an instance of `AcceleratedHNSWParams` instance of `AcceleratedHNSWParams` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:600`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/AcceleratedHNSWParams.java:642`_ _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 aa65c2cf30..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:55`_ +_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:73`_ +_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:204`_ +_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:333`_ +_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:398`_ +_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:480`_ +_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:492`_ +_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:505`_ +_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:547`_ +_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:32`_ +_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 fddfb07773..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,7 +31,7 @@ Multi-layer constructor that supports arbitrary number of layers. | `layerNodes` | the nodes on the layer | | `layerAdjacencies` | adjacency list | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:41`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:56`_ ### getNodesOnLevel @@ -41,7 +41,7 @@ public NodesIterator getNodesOnLevel(int level) Get all nodes on a given level as node 0th ordinals. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:92`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:295`_ ### getNeighbors @@ -62,7 +62,7 @@ Get the neighbors for the node and the level it resides. an instance of NeighborArray -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:110`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:313`_ ### seek @@ -72,7 +72,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Move the pointer to exactly the given level's target. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:135`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:338`_ ### nextNeighbor @@ -82,7 +82,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Iterates over the neighbor list. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:145`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:348`_ ### entryNode @@ -92,7 +92,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns graph's entry point on the top level. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:176`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:379`_ ### maxConn @@ -102,7 +102,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap returns M, the maximum number of connections for a node. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:195`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:398`_ ### neighborCount @@ -112,7 +112,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGrap Returns the neighbor count. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:210`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:413`_ ### size @@ -122,7 +122,7 @@ public int size() Returns the number of nodes in the graph. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:285`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:488`_ ### numLevels @@ -136,7 +136,7 @@ Returns the number of levels in the HNSW graph. the number of levels -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:294`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:497`_ ### dimensions @@ -150,6 +150,6 @@ Gets the vector dimension. the vector dimension -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GPUBuiltHnswGraph.java:303`_ +_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:21`_ +_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 150ef20df2..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:221`_ +_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:363`_ +_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:372`_ +_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:392`_ +_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:402`_ +_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 b9ef604a4a..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:220`_ +_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:305`_ +_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:336`_ +_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:356`_ +_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:365`_ +_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 fc8c021fb0..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:247`_ +_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:330`_ +_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:361`_ +_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:381`_ +_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:390`_ +_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 34aa9402d4..670a67305f 100644 --- a/fern/pages/user_guide/lucene.md +++ b/fern/pages/user_guide/lucene.md @@ -231,7 +231,7 @@ Both `AcceleratedHNSWParams` and `GPUSearchParams` default to a `HEURISTIC` stra - For the accelerated HNSW codecs, set `maxConn` and `beamWidth`, the HNSW parameters you would tune on the CPU. cuVS derives graph degrees and the build algorithm from them. These two values also configure the CPU fallback writer, so one setting covers both paths. - For the GPU search codec, set `buildQuality`. Higher values spend more build time for a higher-quality graph. This codec also passes `graphDegree` into the heuristic, so leave it at its default unless you intend to cap the graph. -Increasing `writerThreads` raises index build concurrency. The accelerated HNSW codecs default to a single writer thread, while the GPU search codec defaults to 32. +`writerThreads` controls native cuVS 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 a5f164b70b..00aa56fd2e 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 @@ -37,6 +37,8 @@ public static enum Strategy { */ 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; @@ -53,6 +55,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; @@ -78,6 +81,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; @@ -95,7 +99,9 @@ public static enum Strategy { /** * Constructs an instance of {@link AcceleratedHNSWParams} with specific parameter values. * - * @param writerThreads Number of cuVS writer threads to use. + * @param writerThreads 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. @@ -112,6 +118,7 @@ public static enum Strategy { */ private AcceleratedHNSWParams( int writerThreads, + int graphThreads, int intermediateGraphDegree, int graphdegree, int hnswLayers, @@ -127,6 +134,7 @@ private AcceleratedHNSWParams( HnswHeuristicType hnswHeuristicType) { super(); this.writerThreads = writerThreads; + this.graphThreads = graphThreads; this.intermediateGraphDegree = intermediateGraphDegree; this.graphdegree = graphdegree; this.hnswLayers = hnswLayers; @@ -143,7 +151,7 @@ private AcceleratedHNSWParams( } /** - * Get the cuVS writer threads parameter + * Get the native cuVS writer threads parameter. * * @return cuVS writer threads parameter */ @@ -151,6 +159,16 @@ 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 * @@ -276,6 +294,8 @@ public HnswHeuristicType getHnswHeuristicType() { public String toString() { return "AcceleratedHNSWParams [writerThreads=" + writerThreads + + ", graphThreads=" + + graphThreads + ", intermediateGraphDegree=" + intermediateGraphDegree + ", graphdegree=" @@ -311,6 +331,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; @@ -326,7 +347,7 @@ public static class Builder { private HnswHeuristicType hnswHeuristicType = DEFAULT_HNSW_HEURISTIC_TYPE; /** - * Set the number of cuVS writer threads while building the index + * 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 +359,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} @@ -521,6 +555,14 @@ private void validate() throws IllegalArgumentException { + MAX_WRITER_THREADS + "]"); } + if (graphThreads < MIN_GRAPH_THREADS || graphThreads > MAX_GRAPH_THREADS) { + throw new IllegalArgumentException( + "graphThreads not in valid range. Valid range: [" + + MIN_GRAPH_THREADS + + ", " + + MAX_GRAPH_THREADS + + "]"); + } if (intermediateGraphDegree < MIN_INT_GRAPH_DEG || intermediateGraphDegree > MAX_INT_GRAPH_DEG) { throw new IllegalArgumentException( @@ -608,6 +650,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 18461bbcfb..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 @@ -20,8 +20,11 @@ import java.util.Random; import java.util.SortedSet; import java.util.TreeSet; +import java.util.concurrent.Callable; import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.store.ByteBuffersDataOutput; +import org.apache.lucene.store.DataOutput; import org.apache.lucene.store.IndexOutput; import org.apache.lucene.util.InfoStream; import org.apache.lucene.util.hnsw.HnswGraph; @@ -82,7 +85,14 @@ public static GPUBuiltHnswGraph createMultiLayerHnswGraph( QuantizationType quantization) throws Throwable { return createMultiLayerHnswGraph( - size, dimensions, adjacencyListMatrix, vectors, hnswLayers, params, quantization); + size, + dimensions, + adjacencyListMatrix, + vectors, + hnswLayers, + params, + quantization, + AcceleratedHNSWParams.DEFAULT_GRAPH_THREADS); } private static GPUBuiltHnswGraph createMultiLayerHnswGraph( @@ -92,7 +102,31 @@ private static GPUBuiltHnswGraph createMultiLayerHnswGraph( List vectors, int hnswLayers, CagraIndexParams params, - QuantizationType quantization) + QuantizationType quantization, + int graphThreads) + throws Throwable { + return createMultiLayerHnswGraph( + size, + dimensions, + adjacencyListMatrix, + vectors, + hnswLayers, + params, + quantization, + graphThreads, + GraphProcessingTrace.disabled()); + } + + private static GPUBuiltHnswGraph createMultiLayerHnswGraph( + int size, + int dimensions, + CuVSMatrix adjacencyListMatrix, + List vectors, + int hnswLayers, + CagraIndexParams params, + QuantizationType quantization, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) throws Throwable { int M = Math.ceilDiv((int) adjacencyListMatrix.columns(), 2); @@ -181,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); + return new GPUBuiltHnswGraph( + size, dimensions, layerNodes, layerAdjacencies, graphThreads, graphProcessingTrace); } catch (Throwable t) { failure = t; throw t; @@ -210,6 +245,46 @@ static GPUBuiltHnswGraph createMultiLayerHnswGraph( CagraIndexParams params, QuantizationType quantization) throws Throwable { + return createMultiLayerHnswGraph( + dimensions, + adjacencyListMatrix, + vectorDataset, + hnswLayers, + params, + quantization, + AcceleratedHNSWParams.DEFAULT_GRAPH_THREADS); + } + + static GPUBuiltHnswGraph createMultiLayerHnswGraph( + int dimensions, + CuVSMatrix adjacencyListMatrix, + CuVSMatrix vectorDataset, + int hnswLayers, + CagraIndexParams params, + QuantizationType quantization, + int graphThreads) + throws Throwable { + return createMultiLayerHnswGraph( + dimensions, + adjacencyListMatrix, + vectorDataset, + hnswLayers, + params, + quantization, + graphThreads, + GraphProcessingTrace.disabled()); + } + + static GPUBuiltHnswGraph createMultiLayerHnswGraph( + int dimensions, + CuVSMatrix adjacencyListMatrix, + CuVSMatrix vectorDataset, + int hnswLayers, + CagraIndexParams params, + QuantizationType quantization, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) + throws Throwable { int size = Math.toIntExact(vectorDataset.size()); // Matrix columns are the stored width: binary vectors are bit-packed, while scalar and float // vectors store one value per dimension. @@ -235,7 +310,15 @@ public int size() { } }; return createMultiLayerHnswGraph( - size, dimensions, adjacencyListMatrix, vectors, hnswLayers, params, quantization); + size, + dimensions, + adjacencyListMatrix, + vectors, + hnswLayers, + params, + quantization, + graphThreads, + graphProcessingTrace); } private static Throwable closeUpperLayerAdjacencies(List layerAdjacencies) { @@ -333,54 +416,162 @@ private static CuVSMatrix buildCagraGraphForSubset( */ public static int[][] writeGraph(GPUBuiltHnswGraph graph, IndexOutput vectorIndex) throws IOException { - // write vectors' neighbors on each level into the vectorIndex file + 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()); + } + + static int[][] writeGraph( + GPUBuiltHnswGraph graph, + IndexOutput vectorIndex, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) + throws IOException { int countOnLevel0 = graph.size(); - int[][] offsets = new int[graph.numLevels()][]; - int[] scratch = new int[graph.maxConn() * 2]; - for (int level = 0; level < graph.numLevels(); level++) { + int numLevels = graph.numLevels(); + int[][] offsets = new int[numLevels][]; + int maxConn = graph.maxConn(); + + int[] level0Nodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(0)); + offsets[0] = new int[level0Nodes.length]; + if (graphThreads > 1 && level0Nodes.length >= GPUBuiltHnswGraph.PARALLEL_MIN_NODES) { + writeLevel0Parallel( + graph, vectorIndex, level0Nodes, offsets[0], countOnLevel0, maxConn, graphThreads); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.SERIALIZATION, + GraphProcessingTrace.Mode.PARALLEL, + GraphProcessingTrace.Reason.ABOVE_THRESHOLD, + graphThreads, + level0Nodes.length); + } else { + writeLevelSerial(graph, vectorIndex, 0, level0Nodes, offsets[0], countOnLevel0, maxConn); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.SERIALIZATION, + GraphProcessingTrace.Mode.SERIAL, + graphThreads <= 1 + ? GraphProcessingTrace.Reason.SINGLE_THREAD + : GraphProcessingTrace.Reason.BELOW_THRESHOLD, + graphThreads, + level0Nodes.length); + } + + for (int level = 1; level < numLevels; level++) { int[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); offsets[level] = new int[sortedNodes.length]; - int nodeOffsetId = 0; - - for (int node : sortedNodes) { - // Get node neighbors - NeighborArray neighbors = graph.getNeighbors(level, node); - // Get the size of the neighbor array - int size = neighbors.size(); - // Write size in VInt as the neighbors list is typically small - long offsetStart = vectorIndex.getFilePointer(); - // Get neighbors - int[] nnodes = neighbors.nodes(); - // Sort them - Arrays.sort(nnodes, 0, size); - // Now that we have sorted, do delta encoding to minimize the required bits to store the - // information - int actualSize = 0; - if (size > 0) { - scratch[0] = nnodes[0]; - actualSize = 1; + writeLevelSerial( + graph, vectorIndex, level, sortedNodes, offsets[level], countOnLevel0, maxConn); + } + return offsets; + } + + /** + * Fixed operational guardrail on nodes processed before task-local buffers are concatenated when + * the encoded-byte limit would otherwise permit a very large wave. Unlike the byte limit below, + * this is a policy cap rather than an encoded-size calculation. + */ + static final int MAX_SERIALIZATION_WAVE_NODES = 1 << 20; + + /** Maximum worst-case encoded payload buffered by one serialization wave. */ + static final long MAX_SERIALIZED_BYTES_PER_WAVE = 64L << 20; + + private static final int MAX_VINT_BYTES = 5; + + private static void writeLevelSerial( + GPUBuiltHnswGraph graph, + IndexOutput out, + int level, + int[] nodes, + int[] offsets, + int countOnLevel0, + int maxConn) + throws IOException { + int[] scratch = new int[maxConn * 2]; + for (int i = 0; i < nodes.length; i++) { + long start = out.getFilePointer(); + encodeNode(graph.getNeighbors(level, nodes[i]), scratch, out, countOnLevel0); + offsets[i] = Math.toIntExact(out.getFilePointer() - start); + } + } + + /** Encodes level zero in bounded waves, then concatenates buffers in node order. */ + private static void writeLevel0Parallel( + GPUBuiltHnswGraph graph, + IndexOutput out, + int[] nodes, + int[] offsets, + int countOnLevel0, + int maxConn, + int graphThreads) + throws IOException { + 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; } - // De-duplication - for (int i = 1; i < size; i++) { - assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; - // Sorting step helps here - if (nnodes[i - 1] == nnodes[i]) { - continue; - } - scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; + 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); } - // Write the size after duplicates are removed - vectorIndex.writeVInt(actualSize); - // Write de-duplicated neighbors - for (int i = 0; i < actualSize; i++) { - vectorIndex.writeVInt(scratch[i]); + } + waveStart = waveEnd; + } + } + + 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.size(); + int actualSize = 0; + if (size > 0) { + int[] nodes = neighbors.nodes(); + Arrays.sort(nodes, 0, size); + scratch[0] = nodes[0]; + actualSize = 1; + for (int i = 1; i < size; i++) { + assert nodes[i] < countOnLevel0 : "node too large: " + nodes[i] + ">=" + countOnLevel0; + if (nodes[i - 1] != nodes[i]) { + scratch[actualSize++] = nodes[i] - nodes[i - 1]; } - offsets[level][nodeOffsetId++] = - Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); } } - // Return offsets (information written while writing the meta info) - return offsets; + out.writeVInt(actualSize); + for (int i = 0; i < actualSize; i++) { + out.writeVInt(scratch[i]); + } } /** 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 b7a50316b4..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 @@ -6,10 +6,16 @@ import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS; +import com.nvidia.cuvs.CuVSDeviceMatrix; +import com.nvidia.cuvs.CuVSHostMatrix; import com.nvidia.cuvs.CuVSMatrix; import com.nvidia.cuvs.RowView; +import java.io.IOException; import java.util.ArrayList; import java.util.List; +import java.util.Optional; +import java.util.concurrent.Callable; +import java.util.function.Supplier; import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.NeighborArray; @@ -31,6 +37,15 @@ 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, + NeighborArray[] layer0Neighbors, + List layerNeighbors) {} + /** * Multi-layer constructor that supports arbitrary number of layers. * @@ -41,50 +56,238 @@ public class GPUBuiltHnswGraph extends HnswGraph { */ public GPUBuiltHnswGraph( int size, int dimensions, List layerNodes, List layerAdjacencies) { + this(size, dimensions, materializeSerial(size, layerNodes, layerAdjacencies)); + } + // Builds a graph while materializing adjacency rows with up to graphThreads threads. + GPUBuiltHnswGraph( + int size, + int dimensions, + List layerNodes, + List layerAdjacencies, + int graphThreads) + throws IOException { + 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) { this.size = size; this.dimensions = dimensions; - this.numLevels = layerAdjacencies.size(); - this.layerNodes = new ArrayList<>(); - this.layerNeighbors = new ArrayList<>(); + this.numLevels = graph.numLevels(); + this.layerNodes = graph.layerNodes(); + this.layerNeighbors = graph.layerNeighbors(); + this.layer0Neighbors = graph.layer0Neighbors(); + } - // Process Layer 0 (base layer with all nodes) - CuVSMatrix layer0Adjacency = layerAdjacencies.get(0); - this.layer0Neighbors = fillNeighborArray(layer0Adjacency, size); + private static MaterializedGraph materializeSerial( + int size, List layerNodes, List layerAdjacencies) { + try { + return materialize(size, layerNodes, layerAdjacencies, 1, GraphProcessingTrace.disabled()); + } catch (IOException impossible) { + throw new AssertionError( + "serial graph materialization cannot fail with IOException", impossible); + } + } - // Process higher layers (1 to numLevels-1) - for (int level = 1; level < numLevels; level++) { + private static MaterializedGraph materialize( + int size, + List layerNodes, + List layerAdjacencies, + int graphThreads, + GraphProcessingTrace graphProcessingTrace) + throws IOException { + List upperLayerNodes = new ArrayList<>(); + List upperLayerNeighbors = new ArrayList<>(); + NeighborArray[] baseLayerNeighbors = + fillNeighborArray(layerAdjacencies.get(0), size, graphThreads, graphProcessingTrace); + + for (int level = 1; level < layerAdjacencies.size(); level++) { int[] nodes = layerNodes.get(level); - CuVSMatrix adjacency = layerAdjacencies.get(level); - this.layerNodes.add(nodes); - this.layerNeighbors.add(fillNeighborArray(adjacency, nodes.length)); + upperLayerNodes.add(nodes); + upperLayerNeighbors.add( + fillNeighborArray( + layerAdjacencies.get(level), nodes.length, graphThreads, graphProcessingTrace)); } + return new MaterializedGraph( + layerAdjacencies.size(), upperLayerNodes, baseLayerNeighbors, upperLayerNeighbors); } /** - * Fills the neighbor array using the adjacency matrix. + * 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 graphThreads maximum threads for this operation, including the caller * @return the NeighborArray */ - private NeighborArray[] fillNeighborArray(CuVSMatrix adjacency, int size) { + private static NeighborArray[] fillNeighborArray( + CuVSMatrix adjacency, int size, int graphThreads, GraphProcessingTrace graphProcessingTrace) + throws IOException { + 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) { + return materializeDeviceAdjacency( + deviceAdjacency, + size, + graphThreads, + GraphCopyMemoryBudget.system(), + () -> newHostMatrix(deviceAdjacency), + graphProcessingTrace); + } + + NeighborArray[] neighbors = fillNeighborArrayParallel(adjacency, size, graphThreads); + graphProcessingTrace.record( + GraphProcessingTrace.Stage.MATERIALIZATION, + GraphProcessingTrace.Mode.PARALLEL, + GraphProcessingTrace.Reason.HOST_SOURCE, + graphThreads, + size); + return neighbors; + } + + private static NeighborArray[] fillNeighborArraySerial(CuVSMatrix source, int size) { NeighborArray[] neighbors = new NeighborArray[size]; - for (int i = 0; i < size; i++) { - RowView rv = adjacency.getRow(i); + fillNeighborRange(source, neighbors, 0, size); + return neighbors; + } + + 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; + } + 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; + } + } + + private static CuVSHostMatrix newHostMatrix(CuVSDeviceMatrix source) { + try (CuVSMatrix.Builder builder = + CuVSMatrix.hostBuilder(source.size(), source.columns(), source.dataType())) { + return builder.build(); + } + } + + static CuVSHostMatrix copyToHost( + CuVSDeviceMatrix source, Supplier hostCopyFactory) { + CuVSHostMatrix hostCopy = hostCopyFactory.get(); + try { + source.toHost(hostCopy); + return hostCopy; + } catch (RuntimeException | Error failure) { + try { + hostCopy.close(); + } catch (RuntimeException | Error closeFailure) { + if (failure != closeFailure) { + failure.addSuppressed(closeFailure); + } + } + throw failure; + } + } + + 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; + } + tasks.add( + () -> { + fillNeighborRange(source, neighbors, start, end); + return null; + }); + } + GraphWorkExecutor.invokeAll(tasks); + return neighbors; + } + + private static void fillNeighborRange( + CuVSMatrix source, NeighborArray[] neighbors, int start, int end) { + for (int i = start; i < end; i++) { + RowView rv = source.getRow(i); if (rv != null && rv.size() > 0) { - neighbors[i] = new NeighborArray((int) rv.size(), true); + NeighborArray nodeNeighbors = new NeighborArray((int) rv.size(), true); for (int j = 0; j < rv.size(); j++) { int neighbor = rv.getAsInt(j); if (neighbor >= 0) { - neighbors[i].addInOrder(neighbor, 1.0f - (j * 0.001f)); + nodeNeighbors.addInOrder(neighbor, 1.0f - (j * 0.001f)); } } + neighbors[i] = nodeNeighbors; } else { neighbors[i] = new NeighborArray(0, true); } } - return neighbors; } /** diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/GraphCopyMemoryBudget.java 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 bbafbe827a..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( @@ -185,9 +187,16 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw dataset, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.NONE); + QuantizationType.NONE, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -272,7 +281,12 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { int dimensions = fieldInfo.getVectorDimension(); GPUBuiltHnswGraph hnswGraph = createSingleVectorHnswGraph(size, dimensions); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + 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 9efcbec3e2..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( @@ -183,10 +185,17 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw dataset, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.BINARY); + QuantizationType.BINARY, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -280,7 +289,12 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + 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 60cc338e4d..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( @@ -210,10 +212,17 @@ private void writeNonTrivialField(FieldInfo fieldInfo, CuVSMatrix dataset) throw dataset, acceleratedHNSWParams.getHnswLayers(), params, - QuantizationType.SCALAR); + QuantizationType.SCALAR, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexOffset = hnswVectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + writeGraph( + hnswGraph, + hnswVectorIndex, + acceleratedHNSWParams.getGraphThreads(), + graphProcessingTrace); long vectorIndexLength = hnswVectorIndex.getFilePointer() - vectorIndexOffset; writeMeta( hnswVectorIndex, @@ -306,7 +315,12 @@ private void writeSingleVectorGraph(FieldInfo fieldInfo) throws IOException { long vectorIndexOffset = hnswVectorIndex.getFilePointer(); // Write the graph to the vector index - int[][] graphLevelNodeOffsets = writeGraph(hnswGraph, hnswVectorIndex); + int[][] graphLevelNodeOffsets = + 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 03b5633a87..0fc0fdaf43 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()); @@ -140,6 +144,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; + } + } +}