Skip to content

[ML] Inference API _services retrieves authorization information directly from EIS #134398

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Sep 18, 2025
Merged

[ML] Inference API _services retrieves authorization information directly from EIS #134398

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
89ef432
get services querying eis gateway for info
jonathan-buttner Sep 9, 2025
ef6e1c9
[CI] Auto commit changes from spotless
Sep 9, 2025
5bf5428
Adding fixes
jonathan-buttner Sep 9, 2025
cecbc82
Merge branch 'ml-eis-services-api-direct' of github.com:jonathan-butt…
jonathan-buttner Sep 9, 2025
1d1adb8
Fixing tests
jonathan-buttner Sep 9, 2025
dfba0f9
Merge branch 'main' of github.com:elastic/elasticsearch into ml-eis-s…
jonathan-buttner Sep 9, 2025
92e982c
Address feedback for javadoc
jonathan-buttner Sep 15, 2025
b86a5e1
Merge branch 'main' into ml-eis-services-api-direct
jonathan-buttner Sep 16, 2025
30cf241
Addressing feedback
jonathan-buttner Sep 18, 2025
17b54a3
Merge branch 'main' of github.com:elastic/elasticsearch into ml-eis-s…
jonathan-buttner Sep 18, 2025
6583ca4
Merge branch 'main' into ml-eis-services-api-direct
jonathan-buttner Sep 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions ...tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import org.elasticsearch.client.Request;
import org.elasticsearch.common.Strings;
import org.elasticsearch.inference.TaskType;
import org.junit.Before;
import org.junit.BeforeClass;

import java.io.IOException;
Expand All @@ -23,6 +24,14 @@

public class InferenceGetServicesIT extends BaseMockEISAuthServerTest {

@Before
public void setUp() throws Exception {
super.setUp();
// Ensure the mock EIS server has an authorized response ready before each test because each test will
// use the services API which makes a call to EIS
mockEISServer.enqueueAuthorizeAllModelsResponse();
Copy link
Member

@davidkyle davidkyle Sep 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The init() method below also calls mockEISServer.enqueueAuthorizeAllModelsResponse(). Is it equivalent to change the annotation on that method to @Before?

Copy link
Contributor Author

@jonathan-buttner jonathan-buttner Sep 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah the original issue for using @BeforeClass is because I ran into some weird issues locally. The original PR that added the @BeforeClass is here: #128640

For some background, when the node for the test starts up it will reach out to EIS and get the auth response. If that fails (there isn't a response queued in the mock server), then the tests will fail. What I was observing is that the base classes static logic would only be executed once regardless of how many subclasses used the base. This resulted in the first test class succeeding but the second test class that leveraged the base would fail. To get around this I added the @BeforeClass and it seemed to fix the issue. The reason we need this in @BeforeClass is because we need a response queued before Elasticsearch is started. Elasticsearch is started only once at the beginning before all the tests run.

}

@BeforeClass
public static void init() {
// Ensure the mock EIS server has an authorized response ready
Expand Down
1 change: 1 addition & 0 deletions .../src/main/java/org/elasticsearch/xpack/inference/mock/TestCompletionServiceExtension.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,7 @@ public static InferenceServiceConfiguration get() {
);

return new InferenceServiceConfiguration.Builder().setService(NAME)
.setName(NAME)
Copy link
Contributor Author

@jonathan-buttner jonathan-buttner Sep 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These aren't really necessary but I was running into a test failure because this field didn't exist until I switch the sorting field. Figured I'd leave the fix in though.

.setTaskTypes(SUPPORTED_TASK_TYPES)
.setConfigurations(configurationMap)
.build();
Expand Down
1 change: 1 addition & 0 deletions .../main/java/org/elasticsearch/xpack/inference/mock/TestDenseInferenceServiceExtension.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,7 @@ public static InferenceServiceConfiguration get() {
);

return new InferenceServiceConfiguration.Builder().setService(NAME)
.setName(NAME)
.setTaskTypes(supportedTaskTypes)
.setConfigurations(configurationMap)
.build();
Expand Down
1 change: 1 addition & 0 deletions ...n/src/main/java/org/elasticsearch/xpack/inference/mock/TestRerankingServiceExtension.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -228,6 +228,7 @@ public static InferenceServiceConfiguration get() {
);

return new InferenceServiceConfiguration.Builder().setService(NAME)
.setName(NAME)
.setTaskTypes(supportedTaskTypes)
.setConfigurations(configurationMap)
.build();
Expand Down
1 change: 1 addition & 0 deletions ...main/java/org/elasticsearch/xpack/inference/mock/TestSparseInferenceServiceExtension.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,7 @@ public static InferenceServiceConfiguration get() {
);

return new InferenceServiceConfiguration.Builder().setService(NAME)
.setName(NAME)
.setTaskTypes(supportedTaskTypes)
.setConfigurations(configurationMap)
.build();
Expand Down
1 change: 1 addition & 0 deletions .../java/org/elasticsearch/xpack/inference/mock/TestStreamingCompletionServiceExtension.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -326,6 +326,7 @@ public static InferenceServiceConfiguration get() {
);

return new InferenceServiceConfiguration.Builder().setService(NAME)
.setName(NAME)
.setTaskTypes(supportedTaskTypes)
.setConfigurations(configurationMap)
.build();
Expand Down
4 changes: 2 additions & 2 deletions ...java/org/elasticsearch/xpack/inference/integration/InferenceRevokeDefaultEndpointsIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,7 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationReturnsEmpty()
try (var service = createElasticInferenceService()) {
ensureAuthorizationCallFinished(service);

assertThat(service.supportedStreamingTasks(), is(EnumSet.noneOf(TaskType.class)));
assertThat(service.supportedStreamingTasks(), is(EnumSet.of(TaskType.CHAT_COMPLETION)));
assertTrue(service.defaultConfigIds().isEmpty());
assertThat(service.supportedTaskTypes(), is(EnumSet.noneOf(TaskType.class)));

Expand Down Expand Up @@ -299,7 +299,7 @@ public void testRemoves_DefaultChatCompletion_V1_WhenAuthorizationDoesNotReturnA
try (var service = createElasticInferenceService()) {
ensureAuthorizationCallFinished(service);

assertThat(service.supportedStreamingTasks(), is(EnumSet.noneOf(TaskType.class)));
assertThat(service.supportedStreamingTasks(), is(EnumSet.of(TaskType.CHAT_COMPLETION)));
assertThat(
service.defaultConfigIds(),
containsInAnyOrder(
Expand Down
3 changes: 3 additions & 0 deletions x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,7 @@
import org.elasticsearch.xpack.inference.external.http.retry.RetrySettings;
import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender;
import org.elasticsearch.xpack.inference.external.http.sender.RequestExecutorServiceSettings;
import org.elasticsearch.xpack.inference.external.http.sender.Sender;
import org.elasticsearch.xpack.inference.highlight.SemanticTextHighlighter;
import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
import org.elasticsearch.xpack.inference.mapper.OffsetSourceFieldMapper;
Expand Down Expand Up @@ -379,6 +380,8 @@ public Collection<?> createComponents(PluginServices services) {
new TransportGetInferenceDiagnosticsAction.ClientManagers(httpClientManager, elasticInferenceServiceHttpClientManager)
);
components.add(inferenceStatsBinding);
components.add(authorizationHandler);
components.add(new PluginComponentBinding<>(Sender.class, elasicInferenceServiceFactory.get().createSender()));
Copy link
Contributor Author

@jonathan-buttner jonathan-buttner Sep 9, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Without this I get a binding error when running


// Only add InferenceServiceNodeLocalRateLimitCalculator (which is a ClusterStateListener) for cluster aware rate limiting,
// if the rate limiting feature flags are enabled, otherwise provide noop implementation
Expand Down
96 changes: 86 additions & 10 deletions ...in/java/org/elasticsearch/xpack/inference/action/TransportGetInferenceServicesAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,36 +7,55 @@

package org.elasticsearch.xpack.inference.action;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.support.ActionFilters;
import org.elasticsearch.action.support.HandledTransportAction;
import org.elasticsearch.action.support.SubscribableListener;
import org.elasticsearch.common.util.concurrent.EsExecutors;
import org.elasticsearch.core.Nullable;
import org.elasticsearch.inference.InferenceService;
import org.elasticsearch.inference.InferenceServiceConfiguration;
import org.elasticsearch.inference.InferenceServiceRegistry;
import org.elasticsearch.inference.TaskType;
import org.elasticsearch.injection.guice.Inject;
import org.elasticsearch.tasks.Task;
import org.elasticsearch.threadpool.ThreadPool;
import org.elasticsearch.transport.TransportService;
import org.elasticsearch.xpack.core.inference.action.GetInferenceServicesAction;
import org.elasticsearch.xpack.inference.external.http.sender.Sender;
import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceService;
import org.elasticsearch.xpack.inference.services.elastic.authorization.ElasticInferenceServiceAuthorizationModel;
import org.elasticsearch.xpack.inference.services.elastic.authorization.ElasticInferenceServiceAuthorizationRequestHandler;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static org.elasticsearch.xpack.inference.InferencePlugin.UTILITY_THREAD_POOL_NAME;

public class TransportGetInferenceServicesAction extends HandledTransportAction<
GetInferenceServicesAction.Request,
GetInferenceServicesAction.Response> {

private static final Logger logger = LogManager.getLogger(TransportGetInferenceServicesAction.class);

private final InferenceServiceRegistry serviceRegistry;
private final ElasticInferenceServiceAuthorizationRequestHandler eisAuthorizationRequestHandler;
private final Sender eisSender;
private final ThreadPool threadPool;

@Inject
public TransportGetInferenceServicesAction(
TransportService transportService,
ActionFilters actionFilters,
InferenceServiceRegistry serviceRegistry
ThreadPool threadPool,
InferenceServiceRegistry serviceRegistry,
ElasticInferenceServiceAuthorizationRequestHandler eisAuthorizationRequestHandler,
Sender sender
) {
super(
GetInferenceServicesAction.NAME,
Expand All @@ -46,6 +65,9 @@ public TransportGetInferenceServicesAction(
EsExecutors.DIRECT_EXECUTOR_SERVICE
);
this.serviceRegistry = serviceRegistry;
this.eisAuthorizationRequestHandler = eisAuthorizationRequestHandler;
this.eisSender = sender;
this.threadPool = threadPool;
}

@Override
Expand All @@ -69,27 +91,81 @@ private void getServiceConfigurationsForTaskType(
.entrySet()
.stream()
.filter(
service -> service.getValue().hideFromConfigurationApi() == false
// exclude EIS here because the hideFromConfigurationApi() is not supported
Copy link
Member

@davidkyle davidkyle Sep 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was slightly confused about "hideFromConfigurationApi() is not supported" in this comment

Suggested change
// exclude EIS here because the hideFromConfigurationApi() is not supported
// Exclude EIS as the EIS specific configurations are handled separately

service -> service.getValue().name().equals(ElasticInferenceService.NAME) == false
&& service.getValue().hideFromConfigurationApi() == false
&& service.getValue().supportedTaskTypes().contains(requestedTaskType)
)
.sorted(Comparator.comparing(service -> service.getValue().name()))
.collect(Collectors.toCollection(ArrayList::new));

getServiceConfigurationsForServices(filteredServices, listener.delegateFailureAndWrap((delegate, configurations) -> {
delegate.onResponse(new GetInferenceServicesAction.Response(configurations));
}));
getServiceConfigurationsForServicesAndEis(listener, filteredServices, requestedTaskType);
}

private void getAllServiceConfigurations(ActionListener<GetInferenceServicesAction.Response> listener) {
var availableServices = serviceRegistry.getServices()
.entrySet()
.stream()
.filter(service -> service.getValue().hideFromConfigurationApi() == false)
.filter(
// exclude EIS here because the hideFromConfigurationApi() is not supported
service -> service.getValue().name().equals(ElasticInferenceService.NAME) == false
&& service.getValue().hideFromConfigurationApi() == false
)
.sorted(Comparator.comparing(service -> service.getValue().name()))
.collect(Collectors.toCollection(ArrayList::new));
getServiceConfigurationsForServices(availableServices, listener.delegateFailureAndWrap((delegate, configurations) -> {
delegate.onResponse(new GetInferenceServicesAction.Response(configurations));
}));

getServiceConfigurationsForServicesAndEis(listener, availableServices, null);
}

private void getServiceConfigurationsForServicesAndEis(
ActionListener<GetInferenceServicesAction.Response> listener,
ArrayList<Map.Entry<String, InferenceService>> availableServices,
@Nullable TaskType requestedTaskType
) {
SubscribableListener.<ElasticInferenceServiceAuthorizationModel>newForked(authModelListener -> {
// Executing on a separate thread because there's a chance the authorization call needs to do some initialization for the Sender
threadPool.executor(UTILITY_THREAD_POOL_NAME).execute(() -> getEisAuthorization(authModelListener, eisSender));
}).<List<InferenceServiceConfiguration>>andThen((configurationListener, authorizationModel) -> {

ActionListener<List<InferenceServiceConfiguration>> mergeEisConfigListener = configurationListener.delegateFailureAndWrap(
(delegate, serviceConfigs) -> {
if (authorizationModel.isAuthorized() == false) {
delegate.onResponse(serviceConfigs);
return;
}

var config = ElasticInferenceService.createConfiguration(authorizationModel.getAuthorizedTaskTypes());
if (requestedTaskType != null && authorizationModel.getAuthorizedTaskTypes().contains(requestedTaskType) == false) {
delegate.onResponse(serviceConfigs);
return;
}

serviceConfigs.add(config);
serviceConfigs.sort(Comparator.comparing(InferenceServiceConfiguration::getService));
delegate.onResponse(serviceConfigs);
}
);

getServiceConfigurationsForServices(availableServices, mergeEisConfigListener);
Copy link
Member

@davidkyle davidkyle Sep 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

getServiceConfigurationsForServices() is a synchronous method and could be written with a return type instead of a listener. I think that would make this code easier to read as you wouldn't need to define the merge listener

}).<List<InferenceServiceConfiguration>>andThen((configurationListener, authorizationModel) -> {

                    var serviceConfigs = getServiceConfigurationsForServices(availableServices);

                    if (authorizationModel.isAuthorized() == false) {
                        delegate.onResponse(serviceConfigs);
                        return;
                    }


                    if (requestedTaskType != null && authorizationModel.getAuthorizedTaskTypes().contains(requestedTaskType) == false) {
                        delegate.onResponse(serviceConfigs);
                        return;
                    }

                    var config = ElasticInferenceService.createConfiguration(authorizationModel.getAuthorizedTaskTypes());
                    serviceConfigs.add(config);
                    serviceConfigs.sort(Comparator.comparing(InferenceServiceConfiguration::getService));
                    delegate.onResponse(serviceConfigs);
                }

Copy link
Contributor Author

@jonathan-buttner jonathan-buttner Sep 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤦‍♂️ thank you, for some reason I thought it needed use a listener.

})
.addListener(
listener.delegateFailureAndWrap(
(delegate, configurations) -> delegate.onResponse(new GetInferenceServicesAction.Response(configurations))
)
);
}

private void getEisAuthorization(ActionListener<ElasticInferenceServiceAuthorizationModel> listener, Sender sender) {
var disabledServiceListener = listener.delegateResponse((delegate, e) -> {
logger.warn(
"Failed to retrieve authorization information from the "
+ "Elastic Inference Service while determining service configurations. Marking service as disabled.",
e
);
delegate.onResponse(ElasticInferenceServiceAuthorizationModel.newDisabledService());
});

eisAuthorizationRequestHandler.getAuthorization(disabledServiceListener, sender);
}

private void getServiceConfigurationsForServices(
Expand All @@ -101,7 +177,7 @@ private void getServiceConfigurationsForServices(
for (var service : services) {
serviceConfigurations.add(service.getValue().getConfiguration());
}
listener.onResponse(serviceConfigurations.stream().toList());
listener.onResponse(serviceConfigurations);
} catch (Exception e) {
listener.onFailure(e);
}
Expand Down
Loading