[8.x] [ML] Inference duration and error metrics (elastic#115876) (elastic#118700)

* [ML] Inference duration and error metrics (elastic#115876)

Add `es.inference.requests.time` metric around `infer` API.

As recommended by OTel spec, errors are determined by the
presence or absence of the `error.type` attribute in the metric.
"error.type" will be the http status code (as a string) if it is
available, otherwise it will be the name of the exception (e.g.
NullPointerException).

Additional notes:
- ApmInferenceStats is merged into InferenceStats. Originally we planned
  to have multiple implementations, but now we're only using APM.
- Request count is now always recorded, even when there are failures
  loading the endpoint configuration.
- Added a hook in streaming for cancel messages, so we can close the
  metrics when a user cancels the stream.

(cherry picked from commit 26870ef)

* fixing switch with class issue

---------

Co-authored-by: Pat Whelan <[email protected]>

Author: 2 people

docs/changelog/115876.yaml

@@ -0,0 +1,5 @@
+pr: 115876
+summary: Inference duration and error metrics
+area: Machine Learning
+type: enhancement
+issues: []

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

@@ -108,7 +108,6 @@
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.IbmWatsonxService;
 import org.elasticsearch.xpack.inference.services.mistral.MistralService;
 import org.elasticsearch.xpack.inference.services.openai.OpenAiService;
-import org.elasticsearch.xpack.inference.telemetry.ApmInferenceStats;
 import org.elasticsearch.xpack.inference.telemetry.InferenceStats;
 
 import java.util.ArrayList;
@@ -256,7 +255,7 @@ public Collection<?> createComponents(PluginServices services) {
         shardBulkInferenceActionFilter.set(actionFilter);
 
         var meterRegistry = services.telemetryProvider().getMeterRegistry();
-        var stats = new PluginComponentBinding<>(InferenceStats.class, ApmInferenceStats.create(meterRegistry));
+        var stats = new PluginComponentBinding<>(InferenceStats.class, InferenceStats.create(meterRegistry));
 
         return List.of(modelRegistry, registry, httpClientManager, stats);
     }

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

@@ -95,6 +95,7 @@ public void request(long n) {
             public void cancel() {
                 if (isClosed.compareAndSet(false, true) && upstream != null) {
                     upstream.cancel();
+                    onCancel();
                 }
             }
         };

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

@@ -7,15 +7,89 @@
 
 package org.elasticsearch.xpack.inference.telemetry;
 
+import org.elasticsearch.ElasticsearchStatusException;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.inference.Model;
+import org.elasticsearch.inference.UnparsedModel;
+import org.elasticsearch.telemetry.metric.LongCounter;
+import org.elasticsearch.telemetry.metric.LongHistogram;
+import org.elasticsearch.telemetry.metric.MeterRegistry;
 
-public interface InferenceStats {
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
-    /**
-     * Increment the counter for a particular value in a thread safe manner.
-     * @param model the model to increment request count for
-     */
-    void incrementRequestCount(Model model);
+import static java.util.Map.entry;
+import static java.util.stream.Stream.concat;
 
-    InferenceStats NOOP = model -> {};
+public record InferenceStats(LongCounter requestCount, LongHistogram inferenceDuration) {
+
+    public InferenceStats {
+        Objects.requireNonNull(requestCount);
+        Objects.requireNonNull(inferenceDuration);
+    }
+
+    public static InferenceStats create(MeterRegistry meterRegistry) {
+        return new InferenceStats(
+            meterRegistry.registerLongCounter(
+                "es.inference.requests.count.total",
+                "Inference API request counts for a particular service, task type, model ID",
+                "operations"
+            ),
+            meterRegistry.registerLongHistogram(
+                "es.inference.requests.time",
+                "Inference API request counts for a particular service, task type, model ID",
+                "ms"
+            )
+        );
+    }
+
+    public static Map<String, Object> modelAttributes(Model model) {
+        return toMap(modelAttributeEntries(model));
+    }
+
+    private static Stream<Map.Entry<String, Object>> modelAttributeEntries(Model model) {
+        var stream = Stream.<Map.Entry<String, Object>>builder()
+            .add(entry("service", model.getConfigurations().getService()))
+            .add(entry("task_type", model.getTaskType().toString()));
+        if (model.getServiceSettings().modelId() != null) {
+            stream.add(entry("model_id", model.getServiceSettings().modelId()));
+        }
+        return stream.build();
+    }
+
+    private static Map<String, Object> toMap(Stream<Map.Entry<String, Object>> stream) {
+        return stream.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+    }
+
+    public static Map<String, Object> responseAttributes(Model model, @Nullable Throwable t) {
+        return toMap(concat(modelAttributeEntries(model), errorAttributes(t)));
+    }
+
+    public static Map<String, Object> responseAttributes(UnparsedModel model, @Nullable Throwable t) {
+        var unknownModelAttributes = Stream.<Map.Entry<String, Object>>builder()
+            .add(entry("service", model.service()))
+            .add(entry("task_type", model.taskType().toString()))
+            .build();
+
+        return toMap(concat(unknownModelAttributes, errorAttributes(t)));
+    }
+
+    public static Map<String, Object> responseAttributes(@Nullable Throwable t) {
+        return toMap(errorAttributes(t));
+    }
+
+    private static Stream<Map.Entry<String, Object>> errorAttributes(@Nullable Throwable t) {
+        if (t == null) {
+            return Stream.of(entry("status_code", 200));
+        } else if (t instanceof ElasticsearchStatusException ese) {
+            return Stream.<Map.Entry<String, Object>>builder()
+                .add(entry("status_code", ese.status().getStatus()))
+                .add(entry("error.type", String.valueOf(ese.status().getStatus())))
+                .build();
+        } else {
+            return Stream.of(entry("error.type", t.getClass().getSimpleName()));
+        }
+    }
 }

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

@@ -0,0 +1,33 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.telemetry;
+
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Objects;
+
+public record InferenceTimer(Instant startTime, Clock clock) {
+
+    public InferenceTimer {
+        Objects.requireNonNull(startTime);
+        Objects.requireNonNull(clock);
+    }
+
+    public static InferenceTimer start() {
+        return start(Clock.systemUTC());
+    }
+
+    public static InferenceTimer start(Clock clock) {
+        return new InferenceTimer(clock.instant(), clock);
+    }
+
+    public long elapsedMillis() {
+        return Duration.between(startTime(), clock().instant()).toMillis();
+    }
+}

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

@@ -1,45 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License
- * 2.0; you may not use this file except in compliance with the Elastic License
- * 2.0.
- */
-
-package org.elasticsearch.xpack.inference.action;
-
-import org.elasticsearch.action.support.ActionFilters;
-import org.elasticsearch.inference.InferenceServiceRegistry;
-import org.elasticsearch.transport.TransportService;
-import org.elasticsearch.xpack.core.inference.action.InferenceAction;
-import org.elasticsearch.xpack.inference.action.task.StreamingTaskManager;
-import org.elasticsearch.xpack.inference.registry.ModelRegistry;
-import org.elasticsearch.xpack.inference.telemetry.InferenceStats;
-
-import static org.mockito.Mockito.mock;
-
-public class TransportInferenceActionTests extends BaseTransportInferenceActionTestCase<InferenceAction.Request> {
-
-    @Override
-    protected BaseTransportInferenceAction<InferenceAction.Request> createAction(
-        TransportService transportService,
-        ActionFilters actionFilters,
-        ModelRegistry modelRegistry,
-        InferenceServiceRegistry serviceRegistry,
-        InferenceStats inferenceStats,
-        StreamingTaskManager streamingTaskManager
-    ) {
-        return new TransportInferenceAction(
-            transportService,
-            actionFilters,
-            modelRegistry,
-            serviceRegistry,
-            inferenceStats,
-            streamingTaskManager
-        );
-    }
-
-    @Override
-    protected InferenceAction.Request createRequest() {
-        return mock();
-    }
-}