diff --git a/x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/xpack/gpu/GPUIndexIT.java b/x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/xpack/gpu/GPUIndexIT.java index 80eadd87cf437..6b7f418004129 100644 --- a/x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/xpack/gpu/GPUIndexIT.java +++ b/x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/xpack/gpu/GPUIndexIT.java @@ -35,9 +35,16 @@ @LuceneTestCase.SuppressCodecs("*") // use our custom codec public class GPUIndexIT extends ESIntegTestCase { + public static class TestGPUPlugin extends GPUPlugin { + @Override + protected boolean isGpuIndexingFeatureAllowed() { + return true; + } + } + @Override protected Collection> nodePlugins() { - return List.of(GPUPlugin.class); + return List.of(TestGPUPlugin.class); } @BeforeClass diff --git a/x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/xpack/gpu/GPUPluginInitializationWithGPUIT.java b/x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/xpack/gpu/GPUPluginInitializationWithGPUIT.java index 721e094352431..62eb1aefe5543 100644 --- a/x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/xpack/gpu/GPUPluginInitializationWithGPUIT.java +++ b/x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/xpack/gpu/GPUPluginInitializationWithGPUIT.java @@ -18,18 +18,20 @@ import org.elasticsearch.indices.IndicesService; import org.elasticsearch.plugins.Plugin; import org.elasticsearch.test.ESIntegTestCase; +import org.junit.After; import java.util.Collection; import java.util.List; import static org.elasticsearch.xpack.gpu.TestVectorsFormatUtils.randomGPUSupportedSimilarity; import static org.hamcrest.Matchers.containsString; +import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.startsWith; public class GPUPluginInitializationWithGPUIT extends ESIntegTestCase { static { - TestCuVSServiceProvider.mockedGPUInfoProvider = SUPPORTEp -> new TestCuVSServiceProvider.TestGPUInfoProvider( + TestCuVSServiceProvider.mockedGPUInfoProvider = p -> new TestCuVSServiceProvider.TestGPUInfoProvider( List.of( new GPUInfo( 0, @@ -44,15 +46,34 @@ public class GPUPluginInitializationWithGPUIT extends ESIntegTestCase { ); } + private static boolean isGpuIndexingFeatureAllowed = true; + + public static class TestGPUPlugin extends GPUPlugin { + + public TestGPUPlugin() { + super(); + } + + @Override + protected boolean isGpuIndexingFeatureAllowed() { + return GPUPluginInitializationWithGPUIT.isGpuIndexingFeatureAllowed; + } + } + + @After + public void reset() { + isGpuIndexingFeatureAllowed = true; + } + @Override protected Collection> nodePlugins() { - return List.of(GPUPlugin.class); + return List.of(TestGPUPlugin.class); } public void testFFOff() { assumeFalse("GPU_FORMAT feature flag disabled", GPUPlugin.GPU_FORMAT.isEnabled()); - GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class); + GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class); VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider(); var format = vectorsFormatProvider.getKnnVectorsFormat(null, null, null); @@ -74,7 +95,7 @@ public void testFFOffIndexSettingNotSupported() { public void testFFOffGPUFormatNull() { assumeFalse("GPU_FORMAT feature flag disabled", GPUPlugin.GPU_FORMAT.isEnabled()); - GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class); + GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class); VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider(); createIndex("index1", Settings.EMPTY); @@ -89,10 +110,10 @@ public void testFFOffGPUFormatNull() { assertNull(format); } - public void testIndexSettingOnIndexTypeSupportedGPUSupported() { + public void testIndexSettingOnIndexAllSupported() { assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled()); - GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class); + GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class); VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider(); createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.TRUE).build()); @@ -110,7 +131,7 @@ public void testIndexSettingOnIndexTypeSupportedGPUSupported() { public void testIndexSettingOnIndexTypeNotSupportedThrows() { assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled()); - GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class); + GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class); VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider(); createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.TRUE).build()); @@ -124,10 +145,31 @@ public void testIndexSettingOnIndexTypeNotSupportedThrows() { assertThat(ex.getMessage(), startsWith("[index.vectors.indexing.use_gpu] doesn't support [index_options.type] of")); } - public void testIndexSettingAutoIndexTypeSupportedGPUSupported() { + public void testIndexSettingOnIndexLicenseNotSupportedThrows() { + assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled()); + isGpuIndexingFeatureAllowed = false; + + GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class); + VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider(); + + createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.TRUE).build()); + IndexSettings settings = getIndexSettings(); + final var indexOptions = DenseVectorFieldTypeTests.randomGpuSupportedIndexOptions(); + + var ex = expectThrows( + IllegalArgumentException.class, + () -> vectorsFormatProvider.getKnnVectorsFormat(settings, indexOptions, randomGPUSupportedSimilarity(indexOptions.getType())) + ); + assertThat( + ex.getMessage(), + equalTo("[index.vectors.indexing.use_gpu] was set to [true], but GPU indexing is a [ENTERPRISE] level feature") + ); + } + + public void testIndexSettingAutoAllSupported() { assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled()); - GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class); + GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class); VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider(); createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.AUTO).build()); @@ -142,10 +184,29 @@ public void testIndexSettingAutoIndexTypeSupportedGPUSupported() { assertNotNull(format); } + public void testIndexSettingAutoLicenseNotSupported() { + assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled()); + isGpuIndexingFeatureAllowed = false; + + GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class); + VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider(); + + createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.AUTO).build()); + IndexSettings settings = getIndexSettings(); + final var indexOptions = DenseVectorFieldTypeTests.randomGpuSupportedIndexOptions(); + + var format = vectorsFormatProvider.getKnnVectorsFormat( + settings, + indexOptions, + randomGPUSupportedSimilarity(indexOptions.getType()) + ); + assertNull(format); + } + public void testIndexSettingAutoIndexTypeNotSupported() { assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled()); - GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class); + GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class); VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider(); createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.AUTO).build()); @@ -163,7 +224,7 @@ public void testIndexSettingAutoIndexTypeNotSupported() { public void testIndexSettingOff() { assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled()); - GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class); + GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class); VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider(); createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.FALSE).build()); diff --git a/x-pack/plugin/gpu/src/main/java/module-info.java b/x-pack/plugin/gpu/src/main/java/module-info.java index 21b2834999249..95655162de6a6 100644 --- a/x-pack/plugin/gpu/src/main/java/module-info.java +++ b/x-pack/plugin/gpu/src/main/java/module-info.java @@ -13,6 +13,7 @@ requires org.elasticsearch.server; requires org.elasticsearch.base; requires org.elasticsearch.gpu; + requires org.elasticsearch.xcore; provides org.elasticsearch.features.FeatureSpecification with org.elasticsearch.xpack.gpu.GPUFeatures; } 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 c33d77b5d57b7..8905ef2581c83 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 @@ -15,8 +15,10 @@ import org.elasticsearch.gpu.codec.ES92GpuHnswVectorsFormat; import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper; import org.elasticsearch.index.mapper.vectors.VectorsFormatProvider; +import org.elasticsearch.license.License; import org.elasticsearch.plugins.Plugin; import org.elasticsearch.plugins.internal.InternalVectorFormatProviderPlugin; +import org.elasticsearch.xpack.core.XPackPlugin; import java.util.List; @@ -24,6 +26,8 @@ public class GPUPlugin extends Plugin implements InternalVectorFormatProviderPlu public static final FeatureFlag GPU_FORMAT = new FeatureFlag("gpu_vectors_indexing"); + private static final License.OperationMode MINIMUM_ALLOWED_LICENSE = License.OperationMode.ENTERPRISE; + /** * An enum for the tri-state value of the `index.vectors.indexing.use_gpu` setting. */ @@ -59,6 +63,12 @@ public List> getSettings() { } } + // Allow tests to override the license state + protected boolean isGpuIndexingFeatureAllowed() { + var licenseState = XPackPlugin.getSharedLicenseState(); + return licenseState != null && licenseState.isAllowedByLicense(MINIMUM_ALLOWED_LICENSE); + } + @Override public VectorsFormatProvider getVectorsFormatProvider() { return (indexSettings, indexOptions, similarity) -> { @@ -75,9 +85,19 @@ public VectorsFormatProvider getVectorsFormatProvider() { "[index.vectors.indexing.use_gpu] was set to [true], but GPU resources are not accessible on the node." ); } + if (isGpuIndexingFeatureAllowed() == false) { + throw new IllegalArgumentException( + "[index.vectors.indexing.use_gpu] was set to [true], but GPU indexing is a [" + + MINIMUM_ALLOWED_LICENSE + + "] level feature" + ); + } return getVectorsFormat(indexOptions, similarity); } - if (gpuMode == GpuMode.AUTO && vectorIndexTypeSupported(indexOptions.getType()) && GPUSupport.isSupported()) { + if (gpuMode == GpuMode.AUTO + && vectorIndexTypeSupported(indexOptions.getType()) + && GPUSupport.isSupported() + && isGpuIndexingFeatureAllowed()) { return getVectorsFormat(indexOptions, similarity); } } diff --git a/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/GPUDenseVectorFieldMapperTests.java b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/GPUDenseVectorFieldMapperTests.java index a97d6fef991d8..79d33ef678c40 100644 --- a/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/GPUDenseVectorFieldMapperTests.java +++ b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/GPUDenseVectorFieldMapperTests.java @@ -34,7 +34,12 @@ public static void setup() { @Override protected Collection getPlugins() { - var plugin = new GPUPlugin(); + var plugin = new GPUPlugin() { + @Override + protected boolean isGpuIndexingFeatureAllowed() { + return true; + } + }; return Collections.singletonList(plugin); }