Skip to content

Commit c40aab2

Browse files
jonathan-buttnerjonathan-buttner
committed
Working on tests
1 parent fe2f1fe commit c40aab2

File tree

5 files changed

+280
-10
lines changed

5 files changed

+280
-10
lines changed

‎x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/InferenceIndicesIT.java‎

Lines changed: 241 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,241 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.inference.integration;
9+
10+
import org.elasticsearch.action.ActionFuture;
11+
import org.elasticsearch.action.search.SearchPhaseExecutionException;
12+
import org.elasticsearch.common.bytes.BytesReference;
13+
import org.elasticsearch.common.settings.Settings;
14+
import org.elasticsearch.inference.InferenceServiceExtension;
15+
import org.elasticsearch.inference.TaskType;
16+
import org.elasticsearch.license.LicenseSettings;
17+
import org.elasticsearch.license.XPackLicenseState;
18+
import org.elasticsearch.plugins.Plugin;
19+
import org.elasticsearch.test.ESIntegTestCase;
20+
import org.elasticsearch.test.ESTestCase;
21+
import org.elasticsearch.xcontent.XContentBuilder;
22+
import org.elasticsearch.xcontent.XContentFactory;
23+
import org.elasticsearch.xcontent.XContentType;
24+
import org.elasticsearch.xpack.core.LocalStateCompositeXPackPlugin;
25+
import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
26+
import org.elasticsearch.xpack.core.inference.action.PutInferenceModelAction;
27+
import org.elasticsearch.xpack.core.ssl.SSLService;
28+
import org.elasticsearch.xpack.inference.InferenceIndex;
29+
import org.elasticsearch.xpack.inference.InferencePlugin;
30+
import org.elasticsearch.xpack.inference.InferenceSecretsIndex;
31+
import org.elasticsearch.xpack.inference.mock.TestDenseInferenceServiceExtension;
32+
import org.elasticsearch.xpack.inference.mock.TestInferenceServicePlugin;
33+
import org.elasticsearch.xpack.inference.mock.TestSparseInferenceServiceExtension;
34+
import org.junit.After;
35+
36+
import java.io.IOException;
37+
import java.nio.file.Path;
38+
import java.util.Collection;
39+
import java.util.HashMap;
40+
import java.util.List;
41+
import java.util.Map;
42+
43+
import static org.hamcrest.CoreMatchers.containsString;
44+
import static org.hamcrest.CoreMatchers.equalTo;
45+
46+
@ESTestCase.WithoutEntitlements // due to dependency issue ES-12435
47+
public class InferenceIndicesIT extends ESIntegTestCase {
48+
49+
private static final String INDEX_ROUTER_ATTRIBUTE = "node.attr.index_router";
50+
private static final String CONFIG_ROUTER = "config";
51+
private static final String SECRETS_ROUTER = "secrets";
52+
53+
private static final Map<String, Object> BBQ_COMPATIBLE_SERVICE_SETTINGS = Map.of(
54+
"model",
55+
"my_model",
56+
"dimensions",
57+
256,
58+
"similarity",
59+
"cosine",
60+
"api_key",
61+
"my_api_key"
62+
);
63+
64+
private final Map<String, TaskType> inferenceIds = new HashMap<>();
65+
66+
public static class LocalStateIndexSettingsInferencePlugin extends LocalStateCompositeXPackPlugin {
67+
private final InferencePlugin inferencePlugin;
68+
69+
public LocalStateIndexSettingsInferencePlugin(final Settings settings, final Path configPath) throws Exception {
70+
super(settings, configPath);
71+
var thisVar = this;
72+
this.inferencePlugin = new InferencePlugin(settings) {
73+
@Override
74+
protected SSLService getSslService() {
75+
return thisVar.getSslService();
76+
}
77+
78+
@Override
79+
protected XPackLicenseState getLicenseState() {
80+
return thisVar.getLicenseState();
81+
}
82+
83+
@Override
84+
public List<InferenceServiceExtension.Factory> getInferenceServiceFactories() {
85+
return List.of(
86+
TestSparseInferenceServiceExtension.TestInferenceService::new,
87+
TestDenseInferenceServiceExtension.TestInferenceService::new
88+
);
89+
}
90+
91+
@Override
92+
public Settings getIndexSettings() {
93+
return InferenceIndex.settingsForTesting(
94+
Settings.builder().put("index.routing.allocation.require.index_router", "config").build()
95+
);
96+
}
97+
98+
@Override
99+
public Settings getSecretsIndexSettings() {
100+
return InferenceSecretsIndex.settingsForTesting(
101+
Settings.builder().put("index.routing.allocation.require.index_router", "secrets").build()
102+
);
103+
}
104+
};
105+
plugins.add(inferencePlugin);
106+
}
107+
108+
}
109+
110+
@After
111+
public void cleanUp() {
112+
// for (var entry : inferenceIds.entrySet()) {
113+
// assertAcked(
114+
// safeGet(
115+
// client().execute(
116+
// DeleteInferenceEndpointAction.INSTANCE,
117+
// new DeleteInferenceEndpointAction.Request(entry.getKey(), entry.getValue(), true, false)
118+
// )
119+
// )
120+
// );
121+
// }
122+
}
123+
124+
@Override
125+
protected Settings nodeSettings(int nodeOrdinal, Settings otherSettings) {
126+
return Settings.builder().put(LicenseSettings.SELF_GENERATED_LICENSE_TYPE.getKey(), "trial").build();
127+
}
128+
129+
@Override
130+
protected Collection<Class<? extends Plugin>> nodePlugins() {
131+
return List.of(LocalStateIndexSettingsInferencePlugin.class, TestInferenceServicePlugin.class);
132+
}
133+
134+
public void testRetrievingInferenceEndpoint_ThrowsException_WhenIndexNodeIsNotAvailable() throws Exception {
135+
final Settings configIndexNodeAttributes = Settings.builder().put(INDEX_ROUTER_ATTRIBUTE, CONFIG_ROUTER).build();
136+
137+
internalCluster().startMasterOnlyNode(configIndexNodeAttributes);
138+
final String configIndexDataNodes = internalCluster().startDataOnlyNode(configIndexNodeAttributes);
139+
140+
internalCluster().startDataOnlyNode(Settings.builder().put(INDEX_ROUTER_ATTRIBUTE, SECRETS_ROUTER).build());
141+
142+
final String inferenceId = "test-index-id";
143+
createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, BBQ_COMPATIBLE_SERVICE_SETTINGS);
144+
145+
// Ensure the inference indices are created and we can retrieve the inference endpoint
146+
GetInferenceModelAction.Request getInferenceEndpointRequest = new GetInferenceModelAction.Request(
147+
inferenceId,
148+
TaskType.TEXT_EMBEDDING,
149+
true
150+
);
151+
var responseFuture = client().execute(GetInferenceModelAction.INSTANCE, getInferenceEndpointRequest);
152+
assertThat(responseFuture.actionGet(TEST_REQUEST_TIMEOUT).getEndpoints().get(0).getInferenceEntityId(), equalTo(inferenceId));
153+
154+
// stop the node that holds the inference index
155+
internalCluster().stopNode(configIndexDataNodes);
156+
157+
var responseFailureFuture = client().execute(GetInferenceModelAction.INSTANCE, getInferenceEndpointRequest);
158+
var exception = expectThrows(SearchPhaseExecutionException.class, () -> responseFailureFuture.actionGet(TEST_REQUEST_TIMEOUT));
159+
160+
assertThat(exception.toString(), containsString("all shards failed"));
161+
assertThat(exception.toString(), containsString("Node not connected"));
162+
assertThat(exception.toString(), containsString(".inference"));
163+
}
164+
165+
public void testRetrievingInferenceEndpoint_ThrowsException_WhenSecretsIndexNodeIsNotAvailable() throws Exception {
166+
final Settings configIndexNodeAttributes = Settings.builder().put(INDEX_ROUTER_ATTRIBUTE, CONFIG_ROUTER).build();
167+
internalCluster().startMasterOnlyNode(configIndexNodeAttributes);
168+
internalCluster().startDataOnlyNode(configIndexNodeAttributes);
169+
170+
String secretIndexDataNodes = internalCluster().startDataOnlyNode(
171+
Settings.builder().put(INDEX_ROUTER_ATTRIBUTE, SECRETS_ROUTER).build()
172+
);
173+
174+
final String inferenceId = "test-secrets-index-id";
175+
createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, BBQ_COMPATIBLE_SERVICE_SETTINGS);
176+
177+
// Ensure the inference indices are created and we can retrieve the inference endpoint
178+
GetInferenceModelAction.Request getInferenceEndpointRequest = new GetInferenceModelAction.Request(
179+
inferenceId,
180+
TaskType.TEXT_EMBEDDING,
181+
true
182+
);
183+
var responseFuture = client().execute(GetInferenceModelAction.INSTANCE, getInferenceEndpointRequest);
184+
assertThat(responseFuture.actionGet(TEST_REQUEST_TIMEOUT).getEndpoints().get(0).getInferenceEntityId(), equalTo(inferenceId));
185+
186+
// stop the node that holds the inference secrets index
187+
internalCluster().stopNode(secretIndexDataNodes);
188+
189+
// We should not be able to create a new inference endpoint because the secrets index is not available
190+
final String inferenceIdFailing = "test-secrets-index-id2";
191+
var responseFailureFuture = createInferenceEndpointAsync(
192+
TaskType.TEXT_EMBEDDING,
193+
inferenceIdFailing,
194+
BBQ_COMPATIBLE_SERVICE_SETTINGS
195+
);
196+
var exception = expectThrows(SearchPhaseExecutionException.class, () -> responseFailureFuture.actionGet(TEST_REQUEST_TIMEOUT));
197+
198+
assertThat(exception.toString(), containsString("all shards failed"));
199+
assertThat(exception.toString(), containsString("Node not connected"));
200+
assertThat(exception.toString(), containsString(".inference-secrets"));
201+
}
202+
203+
private void createInferenceEndpoint(TaskType taskType, String inferenceId, Map<String, Object> serviceSettings) throws IOException {
204+
var responseFuture = createInferenceEndpointAsync(taskType, inferenceId, serviceSettings);
205+
assertThat(responseFuture.actionGet(TEST_REQUEST_TIMEOUT).getModel().getInferenceEntityId(), equalTo(inferenceId));
206+
207+
inferenceIds.put(inferenceId, taskType);
208+
}
209+
210+
private ActionFuture<PutInferenceModelAction.Response> createInferenceEndpointAsync(
211+
TaskType taskType,
212+
String inferenceId,
213+
Map<String, Object> serviceSettings
214+
) throws IOException {
215+
final String service = switch (taskType) {
216+
case TEXT_EMBEDDING -> TestDenseInferenceServiceExtension.TestInferenceService.NAME;
217+
case SPARSE_EMBEDDING -> TestSparseInferenceServiceExtension.TestInferenceService.NAME;
218+
default -> throw new IllegalArgumentException("Unhandled task type [" + taskType + "]");
219+
};
220+
221+
final BytesReference content;
222+
try (XContentBuilder builder = XContentFactory.jsonBuilder()) {
223+
builder.startObject();
224+
builder.field("service", service);
225+
builder.field("service_settings", serviceSettings);
226+
builder.endObject();
227+
228+
content = BytesReference.bytes(builder);
229+
}
230+
231+
PutInferenceModelAction.Request request = new PutInferenceModelAction.Request(
232+
taskType,
233+
inferenceId,
234+
content,
235+
XContentType.JSON,
236+
TEST_REQUEST_TIMEOUT
237+
);
238+
239+
return client().execute(PutInferenceModelAction.INSTANCE, request);
240+
}
241+
}

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceIndex.java‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,10 +30,19 @@ private InferenceIndex() {}
3030
private static final int INDEX_MAPPING_VERSION = 2;
3131

3232
public static Settings settings() {
33-
return Settings.builder()
34-
.put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
35-
.put(IndexMetadata.SETTING_AUTO_EXPAND_REPLICAS, "0-1")
36-
.build();
33+
return builder().build();
34+
}
35+
36+
private static Settings.Builder builder() {
37+
return Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1).put(IndexMetadata.SETTING_AUTO_EXPAND_REPLICAS, "0-1");
38+
}
39+
40+
/**
41+
* This allows tests to set additional settings for the inference index.
42+
* @param extraSettings additional settings to apply to the inference index.
43+
*/
44+
public static Settings settingsForTesting(Settings extraSettings) {
45+
return builder().put(extraSettings).build();
3746
}
3847

3948
/**

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -466,7 +466,7 @@ public Collection<SystemIndexDescriptor> getSystemIndexDescriptors(Settings sett
466466
.setPrimaryIndex(InferenceIndex.INDEX_NAME)
467467
.setDescription("Contains inference service and model configuration")
468468
.setMappings(InferenceIndex.mappings())
469-
.setSettings(InferenceIndex.settings())
469+
.setSettings(getIndexSettings())
470470
.setOrigin(ClientHelper.INFERENCE_ORIGIN)
471471
.setPriorSystemIndexDescriptors(List.of(inferenceIndexV1Descriptor))
472472
.build(),
@@ -476,13 +476,23 @@ public Collection<SystemIndexDescriptor> getSystemIndexDescriptors(Settings sett
476476
.setPrimaryIndex(InferenceSecretsIndex.INDEX_NAME)
477477
.setDescription("Contains inference service secrets")
478478
.setMappings(InferenceSecretsIndex.mappings())
479-
.setSettings(InferenceSecretsIndex.settings())
479+
.setSettings(getSecretsIndexSettings())
480480
.setOrigin(ClientHelper.INFERENCE_ORIGIN)
481481
.setNetNew()
482482
.build()
483483
);
484484
}
485485

486+
// Overridable for tests
487+
protected Settings getIndexSettings() {
488+
return InferenceIndex.settings();
489+
}
490+
491+
// Overridable for tests
492+
protected Settings getSecretsIndexSettings() {
493+
return InferenceSecretsIndex.settings();
494+
}
495+
486496
@Override
487497
public List<ExecutorBuilder<?>> getExecutorBuilders(Settings settingsToUse) {
488498
return List.of(inferenceUtilityExecutor(settings));

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceSecretsIndex.java‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,19 @@ private InferenceSecretsIndex() {}
2929
private static final int INDEX_MAPPING_VERSION = 1;
3030

3131
public static Settings settings() {
32-
return Settings.builder()
33-
.put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
34-
.put(IndexMetadata.SETTING_AUTO_EXPAND_REPLICAS, "0-1")
35-
.build();
32+
return builder().build();
33+
}
34+
35+
private static Settings.Builder builder() {
36+
return Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1).put(IndexMetadata.SETTING_AUTO_EXPAND_REPLICAS, "0-1");
37+
}
38+
39+
/**
40+
* This allows tests to set additional settings for the inference index.
41+
* @param extraSettings additional settings to apply to the inference index.
42+
*/
43+
public static Settings settingsForTesting(Settings extraSettings) {
44+
return builder().put(extraSettings).build();
3645
}
3746

3847
/**

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/registry/ModelRegistry.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,7 @@ public void getModelWithSecrets(String inferenceEntityId, ActionListener<Unparse
268268
SearchRequest modelSearch = client.prepareSearch(InferenceIndex.INDEX_PATTERN, InferenceSecretsIndex.INDEX_PATTERN)
269269
.setQuery(queryBuilder)
270270
.setSize(2)
271+
.setAllowPartialSearchResults(false)
271272
.request();
272273

273274
client.search(modelSearch, searchListener);

0 commit comments

Comments
 (0)