Skip to content

Commit 79d0684

Browse files
ldematteldematte
authored
[GPU] Add license check to the x-pack plugin (#138491) (#138604)
1 parent 8b99a0f commit 79d0684

File tree

5 files changed

+108
-14
lines changed

5 files changed

+108
-14
lines changed

‎x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/xpack/gpu/GPUIndexIT.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,16 @@
3535
@LuceneTestCase.SuppressCodecs("*") // use our custom codec
3636
public class GPUIndexIT extends ESIntegTestCase {
3737

38+
public static class TestGPUPlugin extends GPUPlugin {
39+
@Override
40+
protected boolean isGpuIndexingFeatureAllowed() {
41+
return true;
42+
}
43+
}
44+
3845
@Override
3946
protected Collection<Class<? extends Plugin>> nodePlugins() {
40-
return List.of(GPUPlugin.class);
47+
return List.of(TestGPUPlugin.class);
4148
}
4249

4350
@BeforeClass

‎x-pack/plugin/gpu/src/internalClusterTest/java/org/elasticsearch/xpack/gpu/GPUPluginInitializationWithGPUIT.java‎

Lines changed: 72 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -18,18 +18,20 @@
1818
import org.elasticsearch.indices.IndicesService;
1919
import org.elasticsearch.plugins.Plugin;
2020
import org.elasticsearch.test.ESIntegTestCase;
21+
import org.junit.After;
2122

2223
import java.util.Collection;
2324
import java.util.List;
2425

2526
import static org.elasticsearch.xpack.gpu.TestVectorsFormatUtils.randomGPUSupportedSimilarity;
2627
import static org.hamcrest.Matchers.containsString;
28+
import static org.hamcrest.Matchers.equalTo;
2729
import static org.hamcrest.Matchers.startsWith;
2830

2931
public class GPUPluginInitializationWithGPUIT extends ESIntegTestCase {
3032

3133
static {
32-
TestCuVSServiceProvider.mockedGPUInfoProvider = SUPPORTEp -> new TestCuVSServiceProvider.TestGPUInfoProvider(
34+
TestCuVSServiceProvider.mockedGPUInfoProvider = p -> new TestCuVSServiceProvider.TestGPUInfoProvider(
3335
List.of(
3436
new GPUInfo(
3537
0,
@@ -44,15 +46,34 @@ public class GPUPluginInitializationWithGPUIT extends ESIntegTestCase {
4446
);
4547
}
4648

49+
private static boolean isGpuIndexingFeatureAllowed = true;
50+
51+
public static class TestGPUPlugin extends GPUPlugin {
52+
53+
public TestGPUPlugin() {
54+
super();
55+
}
56+
57+
@Override
58+
protected boolean isGpuIndexingFeatureAllowed() {
59+
return GPUPluginInitializationWithGPUIT.isGpuIndexingFeatureAllowed;
60+
}
61+
}
62+
63+
@After
64+
public void reset() {
65+
isGpuIndexingFeatureAllowed = true;
66+
}
67+
4768
@Override
4869
protected Collection<Class<? extends Plugin>> nodePlugins() {
49-
return List.of(GPUPlugin.class);
70+
return List.of(TestGPUPlugin.class);
5071
}
5172

5273
public void testFFOff() {
5374
assumeFalse("GPU_FORMAT feature flag disabled", GPUPlugin.GPU_FORMAT.isEnabled());
5475

55-
GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class);
76+
GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class);
5677
VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider();
5778

5879
var format = vectorsFormatProvider.getKnnVectorsFormat(null, null, null);
@@ -74,7 +95,7 @@ public void testFFOffIndexSettingNotSupported() {
7495
public void testFFOffGPUFormatNull() {
7596
assumeFalse("GPU_FORMAT feature flag disabled", GPUPlugin.GPU_FORMAT.isEnabled());
7697

77-
GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class);
98+
GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class);
7899
VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider();
79100

80101
createIndex("index1", Settings.EMPTY);
@@ -89,10 +110,10 @@ public void testFFOffGPUFormatNull() {
89110
assertNull(format);
90111
}
91112

92-
public void testIndexSettingOnIndexTypeSupportedGPUSupported() {
113+
public void testIndexSettingOnIndexAllSupported() {
93114
assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled());
94115

95-
GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class);
116+
GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class);
96117
VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider();
97118

98119
createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.TRUE).build());
@@ -110,7 +131,7 @@ public void testIndexSettingOnIndexTypeSupportedGPUSupported() {
110131
public void testIndexSettingOnIndexTypeNotSupportedThrows() {
111132
assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled());
112133

113-
GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class);
134+
GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class);
114135
VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider();
115136

116137
createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.TRUE).build());
@@ -124,10 +145,31 @@ public void testIndexSettingOnIndexTypeNotSupportedThrows() {
124145
assertThat(ex.getMessage(), startsWith("[index.vectors.indexing.use_gpu] doesn't support [index_options.type] of"));
125146
}
126147

127-
public void testIndexSettingAutoIndexTypeSupportedGPUSupported() {
148+
public void testIndexSettingOnIndexLicenseNotSupportedThrows() {
149+
assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled());
150+
isGpuIndexingFeatureAllowed = false;
151+
152+
GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class);
153+
VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider();
154+
155+
createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.TRUE).build());
156+
IndexSettings settings = getIndexSettings();
157+
final var indexOptions = DenseVectorFieldTypeTests.randomGpuSupportedIndexOptions();
158+
159+
var ex = expectThrows(
160+
IllegalArgumentException.class,
161+
() -> vectorsFormatProvider.getKnnVectorsFormat(settings, indexOptions, randomGPUSupportedSimilarity(indexOptions.getType()))
162+
);
163+
assertThat(
164+
ex.getMessage(),
165+
equalTo("[index.vectors.indexing.use_gpu] was set to [true], but GPU indexing is a [ENTERPRISE] level feature")
166+
);
167+
}
168+
169+
public void testIndexSettingAutoAllSupported() {
128170
assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled());
129171

130-
GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class);
172+
GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class);
131173
VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider();
132174

133175
createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.AUTO).build());
@@ -142,10 +184,29 @@ public void testIndexSettingAutoIndexTypeSupportedGPUSupported() {
142184
assertNotNull(format);
143185
}
144186

187+
public void testIndexSettingAutoLicenseNotSupported() {
188+
assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled());
189+
isGpuIndexingFeatureAllowed = false;
190+
191+
GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class);
192+
VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider();
193+
194+
createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.AUTO).build());
195+
IndexSettings settings = getIndexSettings();
196+
final var indexOptions = DenseVectorFieldTypeTests.randomGpuSupportedIndexOptions();
197+
198+
var format = vectorsFormatProvider.getKnnVectorsFormat(
199+
settings,
200+
indexOptions,
201+
randomGPUSupportedSimilarity(indexOptions.getType())
202+
);
203+
assertNull(format);
204+
}
205+
145206
public void testIndexSettingAutoIndexTypeNotSupported() {
146207
assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled());
147208

148-
GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class);
209+
GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class);
149210
VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider();
150211

151212
createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.AUTO).build());
@@ -163,7 +224,7 @@ public void testIndexSettingAutoIndexTypeNotSupported() {
163224
public void testIndexSettingOff() {
164225
assumeTrue("GPU_FORMAT feature flag enabled", GPUPlugin.GPU_FORMAT.isEnabled());
165226

166-
GPUPlugin gpuPlugin = internalCluster().getInstance(GPUPlugin.class);
227+
GPUPlugin gpuPlugin = internalCluster().getInstance(TestGPUPlugin.class);
167228
VectorsFormatProvider vectorsFormatProvider = gpuPlugin.getVectorsFormatProvider();
168229

169230
createIndex("index1", Settings.builder().put(GPUPlugin.VECTORS_INDEXING_USE_GPU_SETTING.getKey(), GPUPlugin.GpuMode.FALSE).build());

‎x-pack/plugin/gpu/src/main/java/module-info.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
requires org.elasticsearch.server;
1414
requires org.elasticsearch.base;
1515
requires org.elasticsearch.gpu;
16+
requires org.elasticsearch.xcore;
1617

1718
provides org.elasticsearch.features.FeatureSpecification with org.elasticsearch.xpack.gpu.GPUFeatures;
1819
}

‎x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/GPUPlugin.java‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,15 +15,19 @@
1515
import org.elasticsearch.gpu.codec.ES92GpuHnswVectorsFormat;
1616
import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
1717
import org.elasticsearch.index.mapper.vectors.VectorsFormatProvider;
18+
import org.elasticsearch.license.License;
1819
import org.elasticsearch.plugins.Plugin;
1920
import org.elasticsearch.plugins.internal.InternalVectorFormatProviderPlugin;
21+
import org.elasticsearch.xpack.core.XPackPlugin;
2022

2123
import java.util.List;
2224

2325
public class GPUPlugin extends Plugin implements InternalVectorFormatProviderPlugin {
2426

2527
public static final FeatureFlag GPU_FORMAT = new FeatureFlag("gpu_vectors_indexing");
2628

29+
private static final License.OperationMode MINIMUM_ALLOWED_LICENSE = License.OperationMode.ENTERPRISE;
30+
2731
/**
2832
* An enum for the tri-state value of the `index.vectors.indexing.use_gpu` setting.
2933
*/
@@ -59,6 +63,12 @@ public List<Setting<?>> getSettings() {
5963
}
6064
}
6165

66+
// Allow tests to override the license state
67+
protected boolean isGpuIndexingFeatureAllowed() {
68+
var licenseState = XPackPlugin.getSharedLicenseState();
69+
return licenseState != null && licenseState.isAllowedByLicense(MINIMUM_ALLOWED_LICENSE);
70+
}
71+
6272
@Override
6373
public VectorsFormatProvider getVectorsFormatProvider() {
6474
return (indexSettings, indexOptions, similarity) -> {
@@ -75,9 +85,19 @@ public VectorsFormatProvider getVectorsFormatProvider() {
7585
"[index.vectors.indexing.use_gpu] was set to [true], but GPU resources are not accessible on the node."
7686
);
7787
}
88+
if (isGpuIndexingFeatureAllowed() == false) {
89+
throw new IllegalArgumentException(
90+
"[index.vectors.indexing.use_gpu] was set to [true], but GPU indexing is a ["
91+
+ MINIMUM_ALLOWED_LICENSE
92+
+ "] level feature"
93+
);
94+
}
7895
return getVectorsFormat(indexOptions, similarity);
7996
}
80-
if (gpuMode == GpuMode.AUTO && vectorIndexTypeSupported(indexOptions.getType()) && GPUSupport.isSupported()) {
97+
if (gpuMode == GpuMode.AUTO
98+
&& vectorIndexTypeSupported(indexOptions.getType())
99+
&& GPUSupport.isSupported()
100+
&& isGpuIndexingFeatureAllowed()) {
81101
return getVectorsFormat(indexOptions, similarity);
82102
}
83103
}

‎x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/GPUDenseVectorFieldMapperTests.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,12 @@ public static void setup() {
3434

3535
@Override
3636
protected Collection<Plugin> getPlugins() {
37-
var plugin = new GPUPlugin();
37+
var plugin = new GPUPlugin() {
38+
@Override
39+
protected boolean isGpuIndexingFeatureAllowed() {
40+
return true;
41+
}
42+
};
3843
return Collections.singletonList(plugin);
3944
}
4045

0 commit comments

Comments
 (0)