Skip to content

Commit 0265d9e

Browse files
[ML] Improve EIS authorization to perform requests on a periodic basis instead of only once (elastic#123639)
* Refactoring * Add internal cluster setting to aid testing * [CI] Auto commit changes from spotless * Allowing the auth interval to be configurable via a setting * Removing unused code * Adding revocation functionality back --------- Co-authored-by: elasticsearchmachine <[email protected]> (cherry picked from commit a88d645) # Conflicts: # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceService.java # x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationHandler.java
1 parent c77c6a6 commit 0265d9e

21 files changed

+1312
-598
lines changed

server/src/main/java/org/elasticsearch/inference/InferenceServiceRegistry.java

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,13 @@ public void init(Client client) {
4242
}
4343

4444
public void onNodeStarted() {
45-
services.values().forEach(InferenceService::onNodeStarted);
45+
for (var service : services.values()) {
46+
try {
47+
service.onNodeStarted();
48+
} catch (Exception e) {
49+
// ignore
50+
}
51+
}
4652
}
4753

4854
public Map<String, InferenceService> getServices() {

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/BaseMockEISAuthServerTest.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@ public class BaseMockEISAuthServerTest extends ESRestTestCase {
3939
.setting("xpack.security.enabled", "true")
4040
// Adding both settings unless one feature flag is disabled in a particular environment
4141
.setting("xpack.inference.elastic.url", mockEISServer::getUrl)
42+
// If we don't disable this there's a very small chance that the authorization code could attempt to make two
43+
// calls which would result in a test failure because the webserver is only expecting a single request
44+
// So to ensure we avoid that all together, this flag indicates that we'll only perform a single authorization request
45+
.setting("xpack.inference.elastic.periodic_authorization_enabled", "false")
4246
// This plugin is located in the inference/qa/test-service-plugin package, look for TestInferenceServicePlugin
4347
.plugin("inference-service-test")
4448
.user("x_pack_rest_user", "x-pack-test-password")

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@
2727
import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
2828
import org.elasticsearch.xpack.inference.registry.ModelRegistry;
2929
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceService;
30-
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceComponents;
31-
import org.elasticsearch.xpack.inference.services.elastic.authorization.ElasticInferenceServiceAuthorizationHandler;
30+
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSettingsTests;
31+
import org.elasticsearch.xpack.inference.services.elastic.authorization.ElasticInferenceServiceAuthorizationRequestHandler;
3232
import org.junit.After;
3333
import org.junit.Before;
3434

@@ -262,7 +262,7 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA
262262

263263
private void ensureAuthorizationCallFinished(ElasticInferenceService service) {
264264
service.onNodeStarted();
265-
service.waitForAuthorizationToComplete(TIMEOUT);
265+
service.waitForFirstAuthorizationToComplete(TIMEOUT);
266266
}
267267

268268
private ElasticInferenceService createElasticInferenceService() {
@@ -272,9 +272,9 @@ private ElasticInferenceService createElasticInferenceService() {
272272
return new ElasticInferenceService(
273273
senderFactory,
274274
createWithEmptySettings(threadPool),
275-
ElasticInferenceServiceComponents.withNoRevokeDelay(gatewayUrl),
275+
ElasticInferenceServiceSettingsTests.create(gatewayUrl),
276276
modelRegistry,
277-
new ElasticInferenceServiceAuthorizationHandler(gatewayUrl, threadPool)
277+
new ElasticInferenceServiceAuthorizationRequestHandler(gatewayUrl, threadPool)
278278
);
279279
}
280280
}

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

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@
118118
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceService;
119119
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceComponents;
120120
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSettings;
121-
import org.elasticsearch.xpack.inference.services.elastic.authorization.ElasticInferenceServiceAuthorizationHandler;
121+
import org.elasticsearch.xpack.inference.services.elastic.authorization.ElasticInferenceServiceAuthorizationRequestHandler;
122122
import org.elasticsearch.xpack.inference.services.elasticsearch.ElasticsearchInternalService;
123123
import org.elasticsearch.xpack.inference.services.googleaistudio.GoogleAiStudioService;
124124
import org.elasticsearch.xpack.inference.services.googlevertexai.GoogleVertexAiService;
@@ -273,14 +273,11 @@ public Collection<?> createComponents(PluginServices services) {
273273
);
274274
elasicInferenceServiceFactory.set(elasticInferenceServiceRequestSenderFactory);
275275

276-
ElasticInferenceServiceSettings inferenceServiceSettings = new ElasticInferenceServiceSettings(settings);
277-
String elasticInferenceUrl = inferenceServiceSettings.getElasticInferenceServiceUrl();
276+
var inferenceServiceSettings = new ElasticInferenceServiceSettings(settings);
277+
inferenceServiceSettings.init(services.clusterService());
278278

279-
var elasticInferenceServiceComponentsInstance = ElasticInferenceServiceComponents.withDefaultRevokeDelay(elasticInferenceUrl);
280-
elasticInferenceServiceComponents.set(elasticInferenceServiceComponentsInstance);
281-
282-
var authorizationHandler = new ElasticInferenceServiceAuthorizationHandler(
283-
elasticInferenceServiceComponentsInstance.elasticInferenceServiceUrl(),
279+
var authorizationHandler = new ElasticInferenceServiceAuthorizationRequestHandler(
280+
inferenceServiceSettings.getElasticInferenceServiceUrl(),
284281
services.threadPool()
285282
);
286283

@@ -289,7 +286,7 @@ public Collection<?> createComponents(PluginServices services) {
289286
context -> new ElasticInferenceService(
290287
elasicInferenceServiceFactory.get(),
291288
serviceComponents.get(),
292-
elasticInferenceServiceComponentsInstance,
289+
inferenceServiceSettings,
293290
modelRegistry,
294291
authorizationHandler
295292
)

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

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -126,11 +126,20 @@ public boolean containsDefaultConfigId(String inferenceEntityId) {
126126
return defaultConfigIds.containsKey(inferenceEntityId);
127127
}
128128

129+
/**
130+
* Adds the default configuration information if it does not already exist internally.
131+
* @param defaultConfigId the default endpoint information
132+
*/
133+
public synchronized void putDefaultIdIfAbsent(InferenceService.DefaultConfigId defaultConfigId) {
134+
defaultConfigIds.putIfAbsent(defaultConfigId.inferenceId(), defaultConfigId);
135+
}
136+
129137
/**
130138
* Set the default inference ids provided by the services
131-
* @param defaultConfigId The default
139+
* @param defaultConfigId The default endpoint information
140+
* @throws IllegalStateException if the {@link InferenceService.DefaultConfigId#inferenceId()} already exists internally
132141
*/
133-
public synchronized void addDefaultIds(InferenceService.DefaultConfigId defaultConfigId) {
142+
public synchronized void addDefaultIds(InferenceService.DefaultConfigId defaultConfigId) throws IllegalStateException {
134143
var config = defaultConfigIds.get(defaultConfigId.inferenceId());
135144
if (config != null) {
136145
throw new IllegalStateException(
Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
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.services.elastic;
9+
10+
import org.elasticsearch.inference.MinimalServiceSettings;
11+
import org.elasticsearch.inference.Model;
12+
13+
public record DefaultModelConfig(Model model, MinimalServiceSettings settings) {
14+
15+
}

0 commit comments

Comments
 (0)