diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/GPUPlugin.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/GPUPlugin.java index 62190bc0fb752..3dcca46d1216b 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/GPUPlugin.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/GPUPlugin.java @@ -12,6 +12,8 @@ import org.elasticsearch.common.util.FeatureFlag; import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper; import org.elasticsearch.index.mapper.vectors.VectorsFormatProvider; +import org.elasticsearch.logging.LogManager; +import org.elasticsearch.logging.Logger; import org.elasticsearch.plugins.Plugin; import org.elasticsearch.plugins.internal.InternalVectorFormatProviderPlugin; import org.elasticsearch.xpack.gpu.codec.ES92GpuHnswSQVectorsFormat; @@ -21,6 +23,8 @@ public class GPUPlugin extends Plugin implements InternalVectorFormatProviderPlugin { + private static final Logger logger = LogManager.getLogger(GPUPlugin.class); + public static final FeatureFlag GPU_FORMAT = new FeatureFlag("gpu_vectors_indexing"); /** @@ -49,6 +53,29 @@ public enum GpuMode { Setting.Property.Dynamic ); + /** The default minimum number of vectors required before building on the GPU. */ + public static final int DEFAULT_MIN_NUM_VECTORS_FOR_GPU_BUILD = 10_000; + + public static final int MIN_NUM_VECTORS_FOR_GPU_BUILD = tinySegmentProperty(); + + public static int tinySegmentProperty() { + int v = DEFAULT_MIN_NUM_VECTORS_FOR_GPU_BUILD; + String str = System.getProperty("gpu.tiny.segment.size"); + if (str != null) { + try { + int parsedValue = Integer.parseInt(str); + if (parsedValue > 1) { + v = parsedValue; + } else { + logger.warn("Ignoring gpu.tiny.segment.size. Value too small:" + parsedValue); + } + } catch (NumberFormatException e) { + logger.warn("Bad gpu.tiny.segment.size. Not a number:" + str); + } + } + return v; + } + @Override public List> getSettings() { if (GPU_FORMAT.isEnabled()) { diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswSQVectorsFormat.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswSQVectorsFormat.java index b62766fb39c3a..44d1da629eea9 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswSQVectorsFormat.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswSQVectorsFormat.java @@ -20,6 +20,7 @@ import java.util.function.Supplier; import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; +import static org.elasticsearch.xpack.gpu.GPUPlugin.MIN_NUM_VECTORS_FOR_GPU_BUILD; import static org.elasticsearch.xpack.gpu.codec.ES92GpuHnswVectorsFormat.DEFAULT_BEAM_WIDTH; import static org.elasticsearch.xpack.gpu.codec.ES92GpuHnswVectorsFormat.DEFAULT_MAX_CONN; @@ -33,18 +34,32 @@ public class ES92GpuHnswSQVectorsFormat extends KnnVectorsFormat { static final int MAXIMUM_BEAM_WIDTH = 3200; private final int maxConn; private final int beamWidth; + // The threshold to use to bypass HNSW graph building for tiny segments on the GPU. + private final int tinySegmentsThreshold; /** The format for storing, reading, merging vectors on disk */ private final FlatVectorsFormat flatVectorsFormat; private final Supplier cuVSResourceManagerSupplier; public ES92GpuHnswSQVectorsFormat() { - this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, null, 7, false); + this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, null, 7, false, CuVSResourceManager::pooling, MIN_NUM_VECTORS_FOR_GPU_BUILD); } public ES92GpuHnswSQVectorsFormat(int maxConn, int beamWidth, Float confidenceInterval, int bits, boolean compress) { + this(maxConn, beamWidth, confidenceInterval, bits, compress, CuVSResourceManager::pooling, MIN_NUM_VECTORS_FOR_GPU_BUILD); + } + + public ES92GpuHnswSQVectorsFormat( + int maxConn, + int beamWidth, + Float confidenceInterval, + int bits, + boolean compress, + Supplier cuVSResourceManagerSupplier, + int tinySegmentsThreshold + ) { super(NAME); - this.cuVSResourceManagerSupplier = CuVSResourceManager::pooling; + this.cuVSResourceManagerSupplier = cuVSResourceManagerSupplier; if (maxConn <= 0 || maxConn > MAXIMUM_MAX_CONN) { throw new IllegalArgumentException( "maxConn must be positive and less than or equal to " + MAXIMUM_MAX_CONN + "; maxConn=" + maxConn @@ -55,8 +70,12 @@ public ES92GpuHnswSQVectorsFormat(int maxConn, int beamWidth, Float confidenceIn "beamWidth must be positive and less than or equal to " + MAXIMUM_BEAM_WIDTH + "; beamWidth=" + beamWidth ); } + if (tinySegmentsThreshold < 2) { + throw new IllegalArgumentException("tinySegmentsThreshold must be greater than 1, got:" + tinySegmentsThreshold); + } this.maxConn = maxConn; this.beamWidth = beamWidth; + this.tinySegmentsThreshold = tinySegmentsThreshold; this.flatVectorsFormat = new ES814ScalarQuantizedVectorsFormat(confidenceInterval, bits, compress); } @@ -67,7 +86,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException state, maxConn, beamWidth, - flatVectorsFormat.fieldsWriter(state) + flatVectorsFormat.fieldsWriter(state), + tinySegmentsThreshold ); } @@ -90,6 +110,8 @@ public String toString() { + maxConn + ", beamWidth=" + beamWidth + + ", tinySegmentsThreshold=" + + tinySegmentsThreshold + ", flatVectorFormat=" + flatVectorsFormat + ")"; diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsFormat.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsFormat.java index 8761b9e12f22a..d1dac2b535383 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsFormat.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsFormat.java @@ -21,6 +21,7 @@ import java.util.function.Supplier; import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.MAX_DIMS_COUNT; +import static org.elasticsearch.xpack.gpu.GPUPlugin.MIN_NUM_VECTORS_FOR_GPU_BUILD; /** * Codec format for GPU-accelerated vector indexes. This format is designed to @@ -38,7 +39,6 @@ public class ES92GpuHnswVectorsFormat extends KnnVectorsFormat { static final int DEFAULT_MAX_CONN = 16; // graph degree public static final int DEFAULT_BEAM_WIDTH = 128; // intermediate graph degree - static final int MIN_NUM_VECTORS_FOR_GPU_BUILD = 2; private static final FlatVectorsFormat flatVectorsFormat = new Lucene99FlatVectorsFormat( FlatVectorScorerUtil.getLucene99FlatVectorsScorer() @@ -49,20 +49,31 @@ public class ES92GpuHnswVectorsFormat extends KnnVectorsFormat { // Intermediate graph degree, the number of connections for each node before pruning private final int beamWidth; private final Supplier cuVSResourceManagerSupplier; + // The threshold to use to bypass HNSW graph building for tiny segments on the GPU. + private final int tinySegmentsThreshold; public ES92GpuHnswVectorsFormat() { - this(CuVSResourceManager::pooling, DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH); + this(DEFAULT_MAX_CONN, DEFAULT_BEAM_WIDTH, CuVSResourceManager::pooling, MIN_NUM_VECTORS_FOR_GPU_BUILD); } public ES92GpuHnswVectorsFormat(int maxConn, int beamWidth) { - this(CuVSResourceManager::pooling, maxConn, beamWidth); - }; + this(maxConn, beamWidth, CuVSResourceManager::pooling, MIN_NUM_VECTORS_FOR_GPU_BUILD); + } - public ES92GpuHnswVectorsFormat(Supplier cuVSResourceManagerSupplier, int maxConn, int beamWidth) { + public ES92GpuHnswVectorsFormat( + int maxConn, + int beamWidth, + Supplier cuVSResourceManagerSupplier, + int tinySegmentsThreshold + ) { super(NAME); + if (tinySegmentsThreshold < 2) { + throw new IllegalArgumentException("tinySegmentsThreshold must be greater than 1, got:" + tinySegmentsThreshold); + } this.cuVSResourceManagerSupplier = cuVSResourceManagerSupplier; this.maxConn = maxConn; this.beamWidth = beamWidth; + this.tinySegmentsThreshold = tinySegmentsThreshold; } @Override @@ -72,7 +83,8 @@ public KnnVectorsWriter fieldsWriter(SegmentWriteState state) throws IOException state, maxConn, beamWidth, - flatVectorsFormat.fieldsWriter(state) + flatVectorsFormat.fieldsWriter(state), + tinySegmentsThreshold ); } @@ -95,6 +107,8 @@ public String toString() { + maxConn + ", beamWidth=" + beamWidth + + ", tinySegmentsThreshold=" + + tinySegmentsThreshold + ", flatVectorFormat=" + flatVectorsFormat.getName() + ")"; diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsWriter.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsWriter.java index 3916afe77caf9..04bcbdc956755 100644 --- a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsWriter.java +++ b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsWriter.java @@ -16,6 +16,8 @@ import org.apache.lucene.codecs.KnnVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatFieldVectorsWriter; import org.apache.lucene.codecs.hnsw.FlatVectorsWriter; +import org.apache.lucene.codecs.lucene95.OffHeapByteVectorValues; +import org.apache.lucene.codecs.lucene95.OffHeapFloatVectorValues; import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter; import org.apache.lucene.index.ByteVectorValues; import org.apache.lucene.index.DocsWithFieldSet; @@ -37,6 +39,10 @@ import org.apache.lucene.util.RamUsageEstimator; import org.apache.lucene.util.hnsw.HnswGraph; import org.apache.lucene.util.hnsw.HnswGraph.NodesIterator; +import org.apache.lucene.util.hnsw.HnswGraphBuilder; +import org.apache.lucene.util.hnsw.NeighborArray; +import org.apache.lucene.util.hnsw.OnHeapHnswGraph; +import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier; import org.apache.lucene.util.packed.DirectMonotonicWriter; import org.apache.lucene.util.quantization.ScalarQuantizer; import org.elasticsearch.core.IOUtils; @@ -61,7 +67,6 @@ import static org.elasticsearch.xpack.gpu.codec.ES92GpuHnswVectorsFormat.LUCENE99_HNSW_VECTOR_INDEX_CODEC_NAME; import static org.elasticsearch.xpack.gpu.codec.ES92GpuHnswVectorsFormat.LUCENE99_HNSW_VECTOR_INDEX_EXTENSION; import static org.elasticsearch.xpack.gpu.codec.ES92GpuHnswVectorsFormat.LUCENE99_VERSION_CURRENT; -import static org.elasticsearch.xpack.gpu.codec.ES92GpuHnswVectorsFormat.MIN_NUM_VECTORS_FOR_GPU_BUILD; /** * Writer that builds an Nvidia Carga Graph on GPU and then writes it into the Lucene99 HNSW format, @@ -78,6 +83,7 @@ final class ES92GpuHnswVectorsWriter extends KnnVectorsWriter { private final int M; private final int beamWidth; private final FlatVectorsWriter flatVectorWriter; + private final int tinySegmentsThreshold; private final List fields = new ArrayList<>(); private boolean finished; @@ -88,13 +94,15 @@ final class ES92GpuHnswVectorsWriter extends KnnVectorsWriter { SegmentWriteState state, int M, int beamWidth, - FlatVectorsWriter flatVectorWriter + FlatVectorsWriter flatVectorWriter, + int tinySegmentsThreshold ) throws IOException { assert cuVSResourceManager != null : "CuVSResources must not be null"; this.cuVSResourceManager = cuVSResourceManager; this.M = M; this.beamWidth = beamWidth; this.flatVectorWriter = flatVectorWriter; + this.tinySegmentsThreshold = tinySegmentsThreshold; if (flatVectorWriter instanceof ES814ScalarQuantizedVectorsFormat.ES814ScalarQuantizedVectorsWriter) { dataType = CuVSMatrix.DataType.BYTE; } else { @@ -177,20 +185,15 @@ private void flushFieldsWithoutMemoryMappedFile(Sorter.DocMap sortMap) throws IO // No tmp file written, or the file cannot be mmapped for (FieldWriter field : fields) { var started = System.nanoTime(); - var fieldInfo = field.fieldInfo; - var numVectors = field.flatFieldVectorsWriter.getVectors().size(); - if (numVectors < MIN_NUM_VECTORS_FOR_GPU_BUILD) { + if (numVectors < tinySegmentsThreshold) { if (logger.isDebugEnabled()) { - logger.debug( - "Skip building carga index; vectors length {} < {} (min for GPU)", - numVectors, - MIN_NUM_VECTORS_FOR_GPU_BUILD - ); + logger.debug("Skip building carga index; vectors length {} < {} (min for GPU)", numVectors, tinySegmentsThreshold); } - // Will not be indexed on the GPU - flushFieldWithMockGraph(fieldInfo, numVectors, sortMap); + // Will not be indexed on the GPU, rather than the CPU + flushFieldBuildingGraphOnCPU(field, sortMap); } else { + var fieldInfo = field.fieldInfo; var cuVSResources = cuVSResourceManager.acquire(numVectors, fieldInfo.getVectorDimension(), CuVSMatrix.DataType.FLOAT); try { var builder = CuVSMatrix.deviceBuilder( @@ -214,12 +217,21 @@ private void flushFieldsWithoutMemoryMappedFile(Sorter.DocMap sortMap) throws IO } } - private void flushFieldWithMockGraph(FieldInfo fieldInfo, int numVectors, Sorter.DocMap sortMap) throws IOException { + void flushFieldBuildingGraphOnCPU(FieldWriter fieldWriter, Sorter.DocMap sortMap) throws IOException { + var fieldInfo = fieldWriter.fieldInfo; + var scorer = flatVectorWriter.getFlatVectorScorer(); + RandomVectorScorerSupplier scorerSupplier = switch (fieldInfo.getVectorEncoding()) { + case BYTE -> throw new AssertionError("bytes not supported"); + case FLOAT32 -> scorer.getRandomVectorScorerSupplier( + fieldInfo.getVectorSimilarityFunction(), + FloatVectorValues.fromFloats(fieldWriter.flatFieldVectorsWriter.getVectors(), fieldInfo.getVectorDimension()) + ); + }; if (sortMap == null) { - generateMockGraphAndWriteMeta(fieldInfo, numVectors); + generateCPUGraphAndWriteMeta(fieldWriter, scorerSupplier); } else { // TODO: use sortMap - generateMockGraphAndWriteMeta(fieldInfo, numVectors); + generateCPUGraphAndWriteMeta(fieldWriter, scorerSupplier); } } @@ -271,7 +283,7 @@ private void generateGpuGraphAndWriteMeta( CuVSMatrix dataset ) throws IOException { try { - assert dataset.size() >= MIN_NUM_VECTORS_FOR_GPU_BUILD; + assert dataset.size() >= tinySegmentsThreshold; long vectorIndexOffset = vectorIndex.getFilePointer(); int[][] graphLevelNodeOffsets = new int[1][]; @@ -289,18 +301,52 @@ private void generateGpuGraphAndWriteMeta( } } - private void generateMockGraphAndWriteMeta(FieldInfo fieldInfo, int datasetSize) throws IOException { - try { - long vectorIndexOffset = vectorIndex.getFilePointer(); - int[][] graphLevelNodeOffsets = new int[1][]; - final HnswGraph graph = writeMockGraph(datasetSize, graphLevelNodeOffsets); - long vectorIndexLength = vectorIndex.getFilePointer() - vectorIndexOffset; - writeMeta(fieldInfo, vectorIndexOffset, vectorIndexLength, datasetSize, graph, graphLevelNodeOffsets); - } catch (IOException e) { - throw e; - } catch (Throwable t) { - throw new IOException("Failed to write GPU index: ", t); + void generateCPUGraphAndWriteMeta(FieldWriter fieldWriter, RandomVectorScorerSupplier scorerSupplier) throws IOException { + final int datasetSize = fieldWriter.flatFieldVectorsWriter.getVectors().size(); + final FieldInfo fieldInfo = fieldWriter.fieldInfo; + OnHeapHnswGraph graph = null; + if (datasetSize > 0) { + graph = buildGraphWithTheCPU(scorerSupplier, datasetSize); + } + writeGraphAndMeta(fieldInfo, graph, datasetSize); + } + + OnHeapHnswGraph buildGraphWithTheCPU(RandomVectorScorerSupplier scorerSupplier, int numVectors) throws IOException { + assert numVectors > 0; + var hnswGraphBuilder = HnswGraphBuilder.create(scorerSupplier, M, beamWidth, HnswGraphBuilder.randSeed); + for (int i = 0; i < numVectors; i++) { + hnswGraphBuilder.addGraphNode(i); + } + return hnswGraphBuilder.getCompletedGraph(); + } + + void writeGraphAndMeta(FieldInfo fieldInfo, OnHeapHnswGraph graph, int datasetSize) throws IOException { + long vectorIndexOffset = vectorIndex.getFilePointer(); + int[][] graphLevelNodeOffsets = writeGraph(graph); // graph may be null + long vectorIndexLength = vectorIndex.getFilePointer() - vectorIndexOffset; + writeMeta(fieldInfo, vectorIndexOffset, vectorIndexLength, datasetSize, graph, graphLevelNodeOffsets); + } + + void createGraphWithCPUAndWriteMeta(FieldInfo fieldInfo, IndexInput input, int size) throws IOException { + OnHeapHnswGraph graph = null; + if (size > 0) { + var vectorsScorer = flatVectorWriter.getFlatVectorScorer(); + var similarity = fieldInfo.getVectorSimilarityFunction(); + var dims = fieldInfo.getVectorDimension(); + final RandomVectorScorerSupplier scoreSupplier = switch (dataType) { + case BYTE -> vectorsScorer.getRandomVectorScorerSupplier( + similarity, + new OffHeapByteVectorValues.DenseOffHeapVectorValues(dims, size, input, dims * Byte.BYTES, vectorsScorer, similarity) + ); + case FLOAT -> vectorsScorer.getRandomVectorScorerSupplier( + similarity, + new OffHeapFloatVectorValues.DenseOffHeapVectorValues(dims, size, input, dims * Float.BYTES, vectorsScorer, similarity) + ); + default -> throw new UnsupportedOperationException("unsupported datatype:" + dataType); + }; + graph = buildGraphWithTheCPU(scoreSupplier, size); } + writeGraphAndMeta(fieldInfo, graph, size); // a null graph is ok } private CagraIndex buildGPUIndex( @@ -371,33 +417,50 @@ private HnswGraph writeGraph(CuVSMatrix cagraGraph, int[][] levelNodeOffsets) th return createMockGraph(maxElementCount, maxGraphDegree); } - // create a mock graph where every node is connected to every other node - private HnswGraph writeMockGraph(int elementCount, int[][] levelNodeOffsets) throws IOException { - if (elementCount == 0) { - return null; - } - int nodeDegree = elementCount - 1; - levelNodeOffsets[0] = new int[elementCount]; - - int[] neighbors = new int[nodeDegree]; - int[] scratch = new int[nodeDegree]; - for (int node = 0; node < elementCount; node++) { - if (nodeDegree > 0) { - for (int j = 0; j < nodeDegree; j++) { - neighbors[j] = j < node ? j : j + 1; // skip self + /** + * Copied from Lucene99HnswVectorsWriter + * @param graph Write the graph in a compressed format + * @return The non-cumulative offsets for the nodes. Should be used to create cumulative offsets. + * @throws IOException if writing to vectorIndex fails + */ + private int[][] writeGraph(OnHeapHnswGraph graph) throws IOException { + if (graph == null) return new int[0][0]; + // write vectors' neighbours on each level into the vectorIndex file + 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[] sortedNodes = NodesIterator.getSortedNodes(graph.getNodesOnLevel(level)); + offsets[level] = new int[sortedNodes.length]; + int nodeOffsetId = 0; + for (int node : sortedNodes) { + NeighborArray neighbors = graph.getNeighbors(level, node); + int size = neighbors.size(); + // Write size in VInt as the neighbors list is typically small + long offsetStart = vectorIndex.getFilePointer(); + int[] nnodes = neighbors.nodes(); + 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; } - scratch[0] = neighbors[0]; - for (int i = 1; i < nodeDegree; i++) { - scratch[i] = neighbors[i] - neighbors[i - 1]; + for (int i = 1; i < size; i++) { + assert nnodes[i] < countOnLevel0 : "node too large: " + nnodes[i] + ">=" + countOnLevel0; + if (nnodes[i - 1] == nnodes[i]) { + continue; + } + scratch[actualSize++] = nnodes[i] - nnodes[i - 1]; } + // Write the size after duplicates are removed + vectorIndex.writeVInt(actualSize); + vectorIndex.writeGroupVInts(scratch, actualSize); + offsets[level][nodeOffsetId++] = Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); } - - long offsetStart = vectorIndex.getFilePointer(); - vectorIndex.writeVInt(nodeDegree); - vectorIndex.writeGroupVInts(scratch, nodeDegree); - levelNodeOffsets[0][node] = Math.toIntExact(vectorIndex.getFilePointer() - offsetStart); } - return createMockGraph(elementCount, nodeDegree); + return offsets; } private static HnswGraph createMockGraph(int elementCount, int graphDegree) { @@ -476,7 +539,7 @@ public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOE try (IndexInput in = mergeState.segmentInfo.dir.openInput(tempRawVectorsFileName, IOContext.DEFAULT)) { var input = FilterIndexInput.unwrapOnlyTest(in); - if (numVectors >= MIN_NUM_VECTORS_FOR_GPU_BUILD) { + if (numVectors > tinySegmentsThreshold) { if (input instanceof MemorySegmentAccessInput memorySegmentAccessInput) { // Direct access to mmapped file final var dataset = DatasetUtils.getInstance() @@ -524,9 +587,7 @@ public void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOE } } } else { - // we don't really need real value for vectors here, - // we just build a mock graph where every node is connected to every other node - generateMockGraphAndWriteMeta(fieldInfo, numVectors); + createGraphWithCPUAndWriteMeta(fieldInfo, input, numVectors); } } catch (Throwable t) { throw new IOException("Failed to merge GPU index: ", t); diff --git a/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/DefaultsTests.java b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/DefaultsTests.java new file mode 100644 index 0000000000000..c20d9227d997a --- /dev/null +++ b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/DefaultsTests.java @@ -0,0 +1,65 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.gpu.codec; + +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.xpack.gpu.GPUPlugin; + +import static org.elasticsearch.xpack.gpu.codec.ES92GpuHnswSQVectorsFormatTests.createES92GpuHnswSQVectorsFormat; +import static org.elasticsearch.xpack.gpu.codec.ES92GpuHnswVectorsFormatTests.createES92GpuHnswVectorsFormat; +import static org.hamcrest.Matchers.startsWith; + +/** Tests for various non-functional things that do not require a GPU or even cuVS. */ +public class DefaultsTests extends ESTestCase { + + public void testDefaultTinySegmentSize() { + assertEquals(GPUPlugin.MIN_NUM_VECTORS_FOR_GPU_BUILD, 10_000); + } + + public void testES92GpuHnswVectorsFormatString() { + var format = createES92GpuHnswVectorsFormat(2); + String expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=2, flatVectorFormat=Lucene99FlatVectorsFormat)"; + assertEquals(expectedStr, format.toString()); + + format = createES92GpuHnswVectorsFormat(10); + expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=10, flatVectorFormat=Lucene99FlatVectorsFormat)"; + assertEquals(expectedStr, format.toString()); + + format = createES92GpuHnswVectorsFormat(1_001_001); + expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=1001001, flatVectorFormat=Lucene99FlatVectorsFormat)"; + assertEquals(expectedStr, format.toString()); + + assertThrows(IllegalArgumentException.class, () -> createES92GpuHnswVectorsFormat(1)); + assertThrows(IllegalArgumentException.class, () -> createES92GpuHnswVectorsFormat(0)); + assertThrows(IllegalArgumentException.class, () -> createES92GpuHnswVectorsFormat(-1)); + } + + public void testES92GpuHnswSQVectorsFormatString() { + var format = createES92GpuHnswSQVectorsFormat(2); + String expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=2, flatVectorFormat=ES814ScalarQuantizedVectorsFormat("; + assertThat(format.toString(), startsWith(expectedStr)); + + format = createES92GpuHnswSQVectorsFormat(10); + expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=10, flatVectorFormat=ES814ScalarQuantizedVectorsFormat("; + assertThat(format.toString(), startsWith(expectedStr)); + + format = createES92GpuHnswSQVectorsFormat(1_001_001); + expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=1001001, flatVectorFormat=ES814ScalarQuantizedVectorsFormat("; + assertThat(format.toString(), startsWith(expectedStr)); + + assertThrows(IllegalArgumentException.class, () -> createES92GpuHnswSQVectorsFormat(1)); + assertThrows(IllegalArgumentException.class, () -> createES92GpuHnswSQVectorsFormat(0)); + assertThrows(IllegalArgumentException.class, () -> createES92GpuHnswSQVectorsFormat(-1)); + } +} diff --git a/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswSQVectorsFormatCPUTests.java b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswSQVectorsFormatCPUTests.java new file mode 100644 index 0000000000000..260d2bdf67a9b --- /dev/null +++ b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswSQVectorsFormatCPUTests.java @@ -0,0 +1,116 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +package org.elasticsearch.xpack.gpu.codec; + +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; +import org.apache.lucene.tests.util.TestUtil; +import org.elasticsearch.common.logging.LogConfigurator; +import org.junit.BeforeClass; + +import static org.hamcrest.Matchers.startsWith; + +/** Runs format tests with a large segment threshold to exercise graph building on the CPU. */ +public class ES92GpuHnswSQVectorsFormatCPUTests extends BaseKnnVectorsFormatTestCase { + + static { + LogConfigurator.loadLog4jPlugins(); + LogConfigurator.configureESLogging(); // native access requires logging to be initialized + } + + static Codec codec; + + /** Format that can only build indices on the CPU, because of the ThrowingCuVSResourceManager. */ + static ES92GpuHnswSQVectorsFormat createES92GpuHnswSQVectorsFormat(int tinySegmentsThreshold) { + return new ES92GpuHnswSQVectorsFormat( + ES92GpuHnswVectorsFormat.DEFAULT_MAX_CONN, + ES92GpuHnswVectorsFormat.DEFAULT_BEAM_WIDTH, + null, + 7, + false, + ThrowingCuVSResourceManager.supplier, + tinySegmentsThreshold + ); + } + + @BeforeClass + public static void beforeClass() { + // Create the format that builds indices on the CPU, because of the tinySegmentThreshold + codec = TestUtil.alwaysKnnVectorsFormat(createES92GpuHnswSQVectorsFormat(Integer.MAX_VALUE)); + } + + @Override + protected Codec getCodec() { + return codec; + } + + public void testKnnVectorsFormatToString() { + KnnVectorsFormat format = createES92GpuHnswSQVectorsFormat(1000000); + String expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=1000000, flatVectorFormat=ES814ScalarQuantizedVectorsFormat"; + assertThat(format.toString(), startsWith(expectedStr)); + + format = createES92GpuHnswSQVectorsFormat(Integer.MAX_VALUE); + expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=2147483647, flatVectorFormat=ES814ScalarQuantizedVectorsFormat"; + assertThat(format.toString(), startsWith(expectedStr)); + + // check the detail values + format = new ES92GpuHnswSQVectorsFormat(); + expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=10000, flatVectorFormat=ES814ScalarQuantizedVectorsFormat"; + assertThat(format.toString(), startsWith(expectedStr)); + + format = new ES92GpuHnswSQVectorsFormat(5, 6, null, 7, false, ThrowingCuVSResourceManager.supplier, 8); + expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=5, beamWidth=6, tinySegmentsThreshold=8, flatVectorFormat=ES814ScalarQuantizedVectorsFormat"; + assertThat(format.toString(), startsWith(expectedStr)); + } + + @Override + protected VectorSimilarityFunction randomSimilarity() { + return VectorSimilarityFunction.values()[random().nextInt(VectorSimilarityFunction.values().length)]; + } + + @Override + protected VectorEncoding randomVectorEncoding() { + return VectorEncoding.FLOAT32; + } + + @Override + public void testRandomBytes() { + // No bytes support + } + + @Override + public void testSortedIndexBytes() { + // No bytes support + } + + @Override + public void testByteVectorScorerIteration() { + // No bytes support + } + + @Override + public void testEmptyByteVectorData() { + // No bytes support + } + + @Override + public void testMergingWithDifferentByteKnnFields() { + // No bytes support + } + + @Override + public void testMismatchedFields() { + // No bytes support + } +} diff --git a/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswSQVectorsFormatTests.java b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswSQVectorsFormatTests.java index f1c13b15795c5..626e57b155c4d 100644 --- a/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswSQVectorsFormatTests.java +++ b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswSQVectorsFormatTests.java @@ -26,10 +26,24 @@ public class ES92GpuHnswSQVectorsFormatTests extends BaseKnnVectorsFormatTestCas static Codec codec; + static ES92GpuHnswSQVectorsFormat createES92GpuHnswSQVectorsFormat(int tinySegmentsThreshold) { + return new ES92GpuHnswSQVectorsFormat( + ES92GpuHnswVectorsFormat.DEFAULT_MAX_CONN, + ES92GpuHnswVectorsFormat.DEFAULT_BEAM_WIDTH, + null, + 7, + false, + CuVSResourceManager::pooling, + tinySegmentsThreshold + ); + } + @BeforeClass public static void beforeClass() { assumeTrue("cuvs not supported", GPUSupport.isSupported(false)); - codec = TestUtil.alwaysKnnVectorsFormat(new ES92GpuHnswSQVectorsFormat()); + + // Create the format that mostly builds indices on the GPU, because of the tinySegmentThreshold + codec = TestUtil.alwaysKnnVectorsFormat(createES92GpuHnswSQVectorsFormat((2))); } @Override diff --git a/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsFormatCPUTests.java b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsFormatCPUTests.java new file mode 100644 index 0000000000000..77e9844e5625a --- /dev/null +++ b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsFormatCPUTests.java @@ -0,0 +1,111 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +package org.elasticsearch.xpack.gpu.codec; + +import org.apache.lucene.codecs.Codec; +import org.apache.lucene.codecs.KnnVectorsFormat; +import org.apache.lucene.index.VectorEncoding; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; +import org.apache.lucene.tests.util.TestUtil; +import org.elasticsearch.common.logging.LogConfigurator; +import org.junit.BeforeClass; + +/** Tests the format while only exercising the CPU-based graph building. */ +// @com.carrotsearch.randomizedtesting.annotations.Repeat(iterations = 100) +public class ES92GpuHnswVectorsFormatCPUTests extends BaseKnnVectorsFormatTestCase { + + static { + LogConfigurator.loadLog4jPlugins(); + LogConfigurator.configureESLogging(); // native access requires logging to be initialized + } + + static Codec codec; + + static ES92GpuHnswVectorsFormat createES92GpuHnswVectorsFormat(int tinySegmentsThreshold) { + return new ES92GpuHnswVectorsFormat( + ES92GpuHnswVectorsFormat.DEFAULT_MAX_CONN, + ES92GpuHnswVectorsFormat.DEFAULT_BEAM_WIDTH, + ThrowingCuVSResourceManager.supplier, + tinySegmentsThreshold + ); + } + + @BeforeClass + public static void beforeClass() { + // Create the format that builds indices on the CPU, because of the tinySegmentThreshold + codec = TestUtil.alwaysKnnVectorsFormat(createES92GpuHnswVectorsFormat(Integer.MAX_VALUE)); + } + + @Override + protected Codec getCodec() { + return codec; + } + + public void testKnnVectorsFormatToString() { + KnnVectorsFormat format = createES92GpuHnswVectorsFormat(1_000_000); + String expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=1000000, flatVectorFormat=Lucene99FlatVectorsFormat)"; + assertEquals(expectedStr, format.toString()); + + format = createES92GpuHnswVectorsFormat(Integer.MAX_VALUE); + expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=2147483647, flatVectorFormat=Lucene99FlatVectorsFormat)"; + assertEquals(expectedStr, format.toString()); + + // check the detail values + format = new ES92GpuHnswVectorsFormat(); + expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=10000, flatVectorFormat=Lucene99FlatVectorsFormat)"; + assertEquals(expectedStr, format.toString()); + + format = new ES92GpuHnswVectorsFormat(5, 6, ThrowingCuVSResourceManager.supplier, 7); + expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " + + "maxConn=5, beamWidth=6, tinySegmentsThreshold=7, flatVectorFormat=Lucene99FlatVectorsFormat)"; + assertEquals(expectedStr, format.toString()); + } + + @Override + protected VectorSimilarityFunction randomSimilarity() { + return VectorSimilarityFunction.values()[random().nextInt(VectorSimilarityFunction.values().length)]; + } + + @Override + protected VectorEncoding randomVectorEncoding() { + return VectorEncoding.FLOAT32; + } + + @Override + public void testRandomBytes() { + // No bytes support + } + + @Override + public void testSortedIndexBytes() { + // No bytes support + } + + @Override + public void testByteVectorScorerIteration() { + // No bytes support + } + + @Override + public void testEmptyByteVectorData() { + // No bytes support + } + + @Override + public void testMergingWithDifferentByteKnnFields() { + // No bytes support + } + + @Override + public void testMismatchedFields() { + // No bytes support + } +} diff --git a/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsFormatTests.java b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsFormatTests.java index e7ce310d15d9b..d6f4101b2c727 100644 --- a/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsFormatTests.java +++ b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ES92GpuHnswVectorsFormatTests.java @@ -27,10 +27,21 @@ public class ES92GpuHnswVectorsFormatTests extends BaseKnnVectorsFormatTestCase static Codec codec; + static ES92GpuHnswVectorsFormat createES92GpuHnswVectorsFormat(int tinySegmentsThreshold) { + return new ES92GpuHnswVectorsFormat( + ES92GpuHnswVectorsFormat.DEFAULT_MAX_CONN, + ES92GpuHnswVectorsFormat.DEFAULT_BEAM_WIDTH, + CuVSResourceManager::pooling, + tinySegmentsThreshold + ); + } + @BeforeClass public static void beforeClass() { assumeTrue("cuvs not supported", GPUSupport.isSupported(false)); - codec = TestUtil.alwaysKnnVectorsFormat(new ES92GpuHnswVectorsFormat()); + + // Create the format that mostly builds indices on the GPU, because of the tinySegmentThreshold + codec = TestUtil.alwaysKnnVectorsFormat(createES92GpuHnswVectorsFormat(2)); } @Override diff --git a/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/GPUDenseVectorFieldMapperTests.java b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/GPUDenseVectorFieldMapperTests.java index 2648691d03eec..74b7684d3cc94 100644 --- a/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/GPUDenseVectorFieldMapperTests.java +++ b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/GPUDenseVectorFieldMapperTests.java @@ -25,6 +25,7 @@ import java.util.Collections; import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.startsWith; public class GPUDenseVectorFieldMapperTests extends DenseVectorFieldMapperTests { @@ -44,7 +45,7 @@ public void testKnnVectorsFormat() throws IOException { // TODO improve test with custom parameters KnnVectorsFormat knnVectorsFormat = getKnnVectorsFormat("hnsw"); String expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " - + "maxConn=16, beamWidth=128, flatVectorFormat=Lucene99FlatVectorsFormat)"; + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=10000, flatVectorFormat=Lucene99FlatVectorsFormat)"; assertEquals(expectedStr, knnVectorsFormat.toString()); } @@ -53,8 +54,8 @@ public void testKnnQuantizedHNSWVectorsFormat() throws IOException { // TOD improve the test with custom parameters KnnVectorsFormat knnVectorsFormat = getKnnVectorsFormat("int8_hnsw"); String expectedStr = "Lucene99HnswVectorsFormat(name=Lucene99HnswVectorsFormat, " - + "maxConn=16, beamWidth=128, flatVectorFormat=ES814ScalarQuantizedVectorsFormat"; - assertTrue(knnVectorsFormat.toString().startsWith(expectedStr)); + + "maxConn=16, beamWidth=128, tinySegmentsThreshold=10000, flatVectorFormat=ES814ScalarQuantizedVectorsFormat"; + assertThat(knnVectorsFormat.toString(), startsWith(expectedStr)); } private KnnVectorsFormat getKnnVectorsFormat(String indexOptionsType) throws IOException { diff --git a/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ThrowingCuVSResourceManager.java b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ThrowingCuVSResourceManager.java new file mode 100644 index 0000000000000..acda1bcbf9149 --- /dev/null +++ b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/ThrowingCuVSResourceManager.java @@ -0,0 +1,38 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.gpu.codec; + +import com.nvidia.cuvs.CuVSMatrix; + +import java.util.function.Supplier; + +/** CuVSResourceManager that always throws. Useful to CPU-only testing. */ +public class ThrowingCuVSResourceManager implements CuVSResourceManager { + + public static Supplier supplier = ThrowingCuVSResourceManager::new; + + @Override + public ManagedCuVSResources acquire(int numVectors, int dims, CuVSMatrix.DataType dataType) { + throw new AssertionError("should not reach here"); + } + + @Override + public void finishedComputation(ManagedCuVSResources resources) { + throw new AssertionError("should not reach here"); + } + + @Override + public void release(ManagedCuVSResources resources) { + throw new AssertionError("should not reach here"); + } + + @Override + public void shutdown() { + throw new AssertionError("should not reach here"); + } +} diff --git a/x-pack/plugin/gpu/src/yamlRestTest/java/org/elasticsearch/xpack/gpu/GPUClientYamlCPUTestSuiteIT.java b/x-pack/plugin/gpu/src/yamlRestTest/java/org/elasticsearch/xpack/gpu/GPUClientYamlCPUTestSuiteIT.java new file mode 100644 index 0000000000000..4a03caac265b1 --- /dev/null +++ b/x-pack/plugin/gpu/src/yamlRestTest/java/org/elasticsearch/xpack/gpu/GPUClientYamlCPUTestSuiteIT.java @@ -0,0 +1,60 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +package org.elasticsearch.xpack.gpu; + +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; + +import org.elasticsearch.test.cluster.ElasticsearchCluster; +import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate; +import org.elasticsearch.test.rest.yaml.ESClientYamlSuiteTestCase; +import org.junit.BeforeClass; +import org.junit.ClassRule; + +/** + * This test sets a large(ish) tiny segment size so that it effectively only exercises code paths + * that build the index on CPU. + */ +public class GPUClientYamlCPUTestSuiteIT extends ESClientYamlSuiteTestCase { + + @BeforeClass + public static void setup() { + assumeTrue("cuvs not supported", GPUSupport.isSupported(false)); + } + + @ClassRule + public static ElasticsearchCluster cluster = createCluster(); + + private static ElasticsearchCluster createCluster() { + var builder = ElasticsearchCluster.local() + .nodes(1) + .module("gpu") + .setting("xpack.license.self_generated.type", "trial") + .setting("xpack.security.enabled", "false") + // set the tiny segment size so that most of the tests exercise CPU index build + .systemProperty("gpu.tiny.segment.size", String.valueOf(Integer.MAX_VALUE)); + + var libraryPath = System.getenv("LD_LIBRARY_PATH"); + if (libraryPath != null) { + builder.environment("LD_LIBRARY_PATH", libraryPath); + } + return builder.build(); + } + + public GPUClientYamlCPUTestSuiteIT(final ClientYamlTestCandidate testCandidate) { + super(testCandidate); + } + + @ParametersFactory + public static Iterable parameters() throws Exception { + return ESClientYamlSuiteTestCase.createParameters(); + } + + @Override + protected String getTestRestCluster() { + return cluster.getHttpAddresses(); + } +} diff --git a/x-pack/plugin/gpu/src/yamlRestTest/java/org/elasticsearch/xpack/gpu/GPUClientYamlTestSuiteIT.java b/x-pack/plugin/gpu/src/yamlRestTest/java/org/elasticsearch/xpack/gpu/GPUClientYamlTestSuiteIT.java index c4e7e936b0111..25863b77c44f3 100644 --- a/x-pack/plugin/gpu/src/yamlRestTest/java/org/elasticsearch/xpack/gpu/GPUClientYamlTestSuiteIT.java +++ b/x-pack/plugin/gpu/src/yamlRestTest/java/org/elasticsearch/xpack/gpu/GPUClientYamlTestSuiteIT.java @@ -29,7 +29,9 @@ private static ElasticsearchCluster createCluster() { .nodes(1) .module("gpu") .setting("xpack.license.self_generated.type", "trial") - .setting("xpack.security.enabled", "false"); + .setting("xpack.security.enabled", "false") + // set the tiny segment size so that most of the tests exercise GPU index build + .systemProperty("gpu.tiny.segment.size", "2"); var libraryPath = System.getenv("LD_LIBRARY_PATH"); if (libraryPath != null) {