diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/InferenceFeatureSetUsage.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/InferenceFeatureSetUsage.java index 00bb8087c9fb3..1a75d079adec2 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/InferenceFeatureSetUsage.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/InferenceFeatureSetUsage.java @@ -20,6 +20,7 @@ import java.io.IOException; import java.util.Collection; +import java.util.List; import java.util.Objects; public class InferenceFeatureSetUsage extends XPackFeatureUsage { @@ -101,6 +102,8 @@ public int hashCode() { } } + public static final InferenceFeatureSetUsage EMPTY = new InferenceFeatureSetUsage(List.of()); + private final Collection modelStats; public InferenceFeatureSetUsage(Collection modelStats) { @@ -129,4 +132,16 @@ public void writeTo(StreamOutput out) throws IOException { public TransportVersion getMinimalSupportedVersion() { return TransportVersions.V_8_12_0; } + + @Override + public boolean equals(Object o) { + if (o == null || getClass() != o.getClass()) return false; + InferenceFeatureSetUsage that = (InferenceFeatureSetUsage) o; + return Objects.equals(modelStats, that.modelStats); + } + + @Override + public int hashCode() { + return Objects.hashCode(modelStats); + } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageAction.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageAction.java index 419d2d3b060a6..7973373b021bb 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageAction.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageAction.java @@ -7,12 +7,15 @@ 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.client.internal.Client; import org.elasticsearch.client.internal.OriginSettingClient; import org.elasticsearch.cluster.ClusterState; import org.elasticsearch.cluster.service.ClusterService; +import org.elasticsearch.common.Strings; import org.elasticsearch.inference.ModelConfigurations; import org.elasticsearch.inference.TaskType; import org.elasticsearch.injection.guice.Inject; @@ -33,6 +36,8 @@ public class TransportInferenceUsageAction extends XPackUsageFeatureTransportAction { + private final Logger logger = LogManager.getLogger(TransportInferenceUsageAction.class); + private final Client client; @Inject @@ -55,7 +60,7 @@ protected void localClusterStateOperation( ActionListener listener ) { GetInferenceModelAction.Request getInferenceModelAction = new GetInferenceModelAction.Request("_all", TaskType.ANY, false); - client.execute(GetInferenceModelAction.INSTANCE, getInferenceModelAction, listener.delegateFailureAndWrap((delegate, response) -> { + client.execute(GetInferenceModelAction.INSTANCE, getInferenceModelAction, ActionListener.wrap(response -> { Map stats = new TreeMap<>(); for (ModelConfigurations model : response.getEndpoints()) { String statKey = model.getService() + ":" + model.getTaskType().name(); @@ -66,7 +71,10 @@ protected void localClusterStateOperation( stat.add(); } InferenceFeatureSetUsage usage = new InferenceFeatureSetUsage(stats.values()); - delegate.onResponse(new XPackUsageFeatureResponse(usage)); + listener.onResponse(new XPackUsageFeatureResponse(usage)); + }, e -> { + logger.warn(Strings.format("Retrieving inference usage failed with error: %s", e.getMessage()), e); + listener.onResponse(new XPackUsageFeatureResponse(InferenceFeatureSetUsage.EMPTY)); })); } } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageActionTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageActionTests.java index 8ffd421959f47..4fe6cd7995850 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageActionTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/action/TransportInferenceUsageActionTests.java @@ -38,6 +38,7 @@ import java.util.List; +import static org.elasticsearch.xpack.inference.Utils.TIMEOUT; import static org.hamcrest.Matchers.hasSize; import static org.hamcrest.core.Is.is; import static org.mockito.ArgumentMatchers.any; @@ -116,4 +117,19 @@ public void test() throws Exception { assertThat(source.getValue("models.2.task_type"), is("TEXT_EMBEDDING")); assertThat(source.getValue("models.2.count"), is(3)); } + + public void testFailureReturnsEmptyUsage() { + doAnswer(invocation -> { + ActionListener listener = invocation.getArgument(2); + listener.onFailure(new IllegalArgumentException("invalid field")); + return Void.TYPE; + }).when(client).execute(any(GetInferenceModelAction.class), any(), any()); + + var future = new PlainActionFuture(); + action.localClusterStateOperation(mock(Task.class), mock(XPackUsageRequest.class), mock(ClusterState.class), future); + + var usage = future.actionGet(TIMEOUT); + var inferenceUsage = (InferenceFeatureSetUsage) usage.getUsage(); + assertThat(inferenceUsage, is(InferenceFeatureSetUsage.EMPTY)); + } }