[ML] Return empty usage statistics in the event of an failure (elastic#124799)

* Returning empty usage on failure

* Adding logging

Author: jonathan-buttner

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> modelStats;
 
     public InferenceFeatureSetUsage(Collection<ModelStats> 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);
+    }
 }

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<XPackUsageFeatureResponse> 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<String, InferenceFeatureSetUsage.ModelStats> 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));
         }));
     }
 }

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<GetInferenceModelAction.Response> 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<XPackUsageFeatureResponse>();
+        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));
+    }
 }