Skip to content

Commit a88d645

Browse files
jonathan-buttnerelasticsearchmachine
andauthored
[ML] Improve EIS authorization to perform requests on a periodic basis instead of only once (#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]>
1 parent 4841f43 commit a88d645

20 files changed

+1298
-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

@@ -270,7 +270,7 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA
270270

271271
private void ensureAuthorizationCallFinished(ElasticInferenceService service) {
272272
service.onNodeStarted();
273-
service.waitForAuthorizationToComplete(TIMEOUT);
273+
service.waitForFirstAuthorizationToComplete(TIMEOUT);
274274
}
275275

276276
private ElasticInferenceService createElasticInferenceService() {
@@ -280,9 +280,9 @@ private ElasticInferenceService createElasticInferenceService() {
280280
return new ElasticInferenceService(
281281
senderFactory,
282282
createWithEmptySettings(threadPool),
283-
ElasticInferenceServiceComponents.withNoRevokeDelay(gatewayUrl),
283+
ElasticInferenceServiceSettingsTests.create(gatewayUrl),
284284
modelRegistry,
285-
new ElasticInferenceServiceAuthorizationHandler(gatewayUrl, threadPool)
285+
new ElasticInferenceServiceAuthorizationRequestHandler(gatewayUrl, threadPool)
286286
);
287287
}
288288
}

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;
@@ -274,14 +274,11 @@ public Collection<?> createComponents(PluginServices services) {
274274
);
275275
elasicInferenceServiceFactory.set(elasticInferenceServiceRequestSenderFactory);
276276

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

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

@@ -290,7 +287,7 @@ public Collection<?> createComponents(PluginServices services) {
290287
context -> new ElasticInferenceService(
291288
elasicInferenceServiceFactory.get(),
292289
serviceComponents.get(),
293-
elasticInferenceServiceComponentsInstance,
290+
inferenceServiceSettings,
294291
modelRegistry,
295292
authorizationHandler
296293
)
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)