Merge remote-tracking branch 'origin/main' into write_coordination_pool

Author: Tim-Brooks

docs/changelog/129161.yaml

@@ -0,0 +1,5 @@
+pr: 129161
+summary: Add Telemetry for models without adaptive allocations
+area: Machine Learning
+type: enhancement
+issues: []

muted-tests.yml

@@ -532,6 +532,12 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.single_node.PushQueriesIT
   method: testCaseInsensitiveEquality {KEYWORD}
   issue: https://github.com/elastic/elasticsearch/issues/129422
+- class: org.elasticsearch.xpack.esql.qa.multi_node.EsqlSpecIT
+  method: test {knn-function.KnnSearchWithKOption ASYNC}
+  issue: https://github.com/elastic/elasticsearch/issues/129447
+- class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeIT
+  method: test
+  issue: https://github.com/elastic/elasticsearch/issues/129453
 
 # Examples:
 #

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

@@ -212,7 +212,6 @@ public static class Builder extends FieldMapper.Builder {
 
         private final Parameter<Map<String, String>> meta = Parameter.metaParam();
 
-        private MinimalServiceSettings resolvedModelSettings;
         private Function<MapperBuilderContext, ObjectMapper> inferenceFieldBuilder;
 
         public static Builder from(SemanticTextFieldMapper mapper) {
@@ -235,14 +234,18 @@ public Builder(
             super(name);
             this.modelRegistry = modelRegistry;
             this.useLegacyFormat = InferenceMetadataFieldsMapper.isEnabled(indexSettings.getSettings()) == false;
-            this.inferenceFieldBuilder = c -> createInferenceField(
-                c,
-                indexSettings.getIndexVersionCreated(),
-                useLegacyFormat,
-                resolvedModelSettings,
-                bitSetProducer,
-                indexSettings
-            );
+            this.inferenceFieldBuilder = c -> {
+                // Resolve the model setting from the registry if it has not been set yet.
+                var resolvedModelSettings = modelSettings.get() != null ? modelSettings.get() : getResolvedModelSettings(c, false);
+                return createInferenceField(
+                    c,
+                    indexSettings.getIndexVersionCreated(),
+                    useLegacyFormat,
+                    resolvedModelSettings,
+                    bitSetProducer,
+                    indexSettings
+                );
+            };
         }
 
         public Builder setInferenceId(String id) {
@@ -283,26 +286,26 @@ protected void merge(FieldMapper mergeWith, Conflicts conflicts, MapperMergeCont
             inferenceFieldBuilder = c -> mergedInferenceField;
         }
 
-        @Override
-        public SemanticTextFieldMapper build(MapperBuilderContext context) {
-            if (useLegacyFormat && copyTo.copyToFields().isEmpty() == false) {
-                throw new IllegalArgumentException(CONTENT_TYPE + " field [" + leafName() + "] does not support [copy_to]");
-            }
-            if (useLegacyFormat && multiFieldsBuilder.hasMultiFields()) {
-                throw new IllegalArgumentException(CONTENT_TYPE + " field [" + leafName() + "] does not support multi-fields");
+        /**
+         * Returns the {@link MinimalServiceSettings} defined in this builder if set;
+         * otherwise, resolves and returns the settings from the registry.
+         */
+        private MinimalServiceSettings getResolvedModelSettings(MapperBuilderContext context, boolean logWarning) {
+            if (context.getMergeReason() == MapperService.MergeReason.MAPPING_RECOVERY) {
+                // the model registry is not available yet
+                return null;
             }
-
-            if (context.getMergeReason() != MapperService.MergeReason.MAPPING_RECOVERY && modelSettings.get() == null) {
-                try {
-                    /*
-                     * If the model is not already set and we are not in a recovery scenario, resolve it using the registry.
-                     * Note: We do not set the model in the mapping at this stage. Instead, the model will be added through
-                     * a mapping update during the first ingestion.
-                     * This approach allows mappings to reference inference endpoints that may not yet exist.
-                     * The only requirement is that the referenced inference endpoint must be available at the time of ingestion.
-                     */
-                    resolvedModelSettings = modelRegistry.getMinimalServiceSettings(inferenceId.get());
-                } catch (ResourceNotFoundException exc) {
+            try {
+                /*
+                 * If the model is not already set and we are not in a recovery scenario, resolve it using the registry.
+                 * Note: We do not set the model in the mapping at this stage. Instead, the model will be added through
+                 * a mapping update during the first ingestion.
+                 * This approach allows mappings to reference inference endpoints that may not yet exist.
+                 * The only requirement is that the referenced inference endpoint must be available at the time of ingestion.
+                 */
+                return modelRegistry.getMinimalServiceSettings(inferenceId.get());
+            } catch (ResourceNotFoundException exc) {
+                if (logWarning) {
                     /* We allow the inference ID to be unregistered at this point.
                      * This will delay the creation of sub-fields, so indexing and querying for this field won't work
                      * until the corresponding inference endpoint is created.
@@ -315,8 +318,22 @@ public SemanticTextFieldMapper build(MapperBuilderContext context) {
                         inferenceId.get()
                     );
                 }
-            } else {
-                resolvedModelSettings = modelSettings.get();
+                return null;
+            }
+        }
+
+        @Override
+        public SemanticTextFieldMapper build(MapperBuilderContext context) {
+            if (useLegacyFormat && copyTo.copyToFields().isEmpty() == false) {
+                throw new IllegalArgumentException(CONTENT_TYPE + " field [" + leafName() + "] does not support [copy_to]");
+            }
+            if (useLegacyFormat && multiFieldsBuilder.hasMultiFields()) {
+                throw new IllegalArgumentException(CONTENT_TYPE + " field [" + leafName() + "] does not support multi-fields");
+            }
+
+            var resolvedModelSettings = modelSettings.get();
+            if (modelSettings.get() == null) {
+                resolvedModelSettings = getResolvedModelSettings(context, true);
             }
 
             if (modelSettings.get() != null) {

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MlMetrics.java

@@ -272,6 +272,25 @@ private void registerMasterNodeMetrics(MeterRegistry meterRegistry) {
                 () -> new LongWithAttributes(trainedModelAllocationCounts.trainedModelsFailedAllocations, isMasterMap)
             )
         );
+        metrics.add(
+            meterRegistry.registerLongGauge(
+                "es.ml.trained_models.deployment.fixed_allocations.current",
+                "Sum of current trained model allocations that do not use adaptive allocations (either enabled or disabled)",
+                "allocations",
+                () -> new LongWithAttributes(trainedModelAllocationCounts.deploymentsWithFixedAllocations, isMasterMap)
+            )
+        );
+        /*
+         * AdaptiveAllocationsScalerService tracks the number of allocations with adaptive allocations enabled.
+         */
+        metrics.add(
+            meterRegistry.registerLongGauge(
+                "es.ml.trained_models.deployment.disabled_adaptive_allocations.current",
+                "Sum of current trained model allocations that have adaptive allocations disabled",
+                "allocations",
+                () -> new LongWithAttributes(trainedModelAllocationCounts.deploymentsWithDisabledAdaptiveAllocations, isMasterMap)
+            )
+        );
     }
 
     @Override
@@ -484,17 +503,28 @@ static TrainedModelAllocationCounts findTrainedModelAllocationCounts(TrainedMode
         int trainedModelsTargetAllocations = 0;
         int trainedModelsCurrentAllocations = 0;
         int trainedModelsFailedAllocations = 0;
+        int deploymentsWithFixedAllocations = 0;
+        int deploymentsWithDisabledAdaptiveAllocations = 0;
 
         for (TrainedModelAssignment trainedModelAssignment : metadata.allAssignments().values()) {
             trainedModelsTargetAllocations += trainedModelAssignment.totalTargetAllocations();
-            trainedModelsCurrentAllocations += trainedModelAssignment.totalCurrentAllocations();
             trainedModelsFailedAllocations += trainedModelAssignment.totalFailedAllocations();
+            trainedModelsCurrentAllocations += trainedModelAssignment.totalCurrentAllocations();
+
+            if (trainedModelAssignment.getAdaptiveAllocationsSettings() == null) {
+                deploymentsWithFixedAllocations += 1;
+            } else if ((trainedModelAssignment.getAdaptiveAllocationsSettings().getEnabled() == null)
+                || (trainedModelAssignment.getAdaptiveAllocationsSettings().getEnabled() == false)) {
+                    deploymentsWithDisabledAdaptiveAllocations += 1;
+                }
         }
 
         return new TrainedModelAllocationCounts(
             trainedModelsTargetAllocations,
             trainedModelsCurrentAllocations,
-            trainedModelsFailedAllocations
+            trainedModelsFailedAllocations,
+            deploymentsWithFixedAllocations,
+            deploymentsWithDisabledAdaptiveAllocations
         );
     }
 
@@ -556,8 +586,10 @@ record MlTaskStatusCounts(
     record TrainedModelAllocationCounts(
         int trainedModelsTargetAllocations,
         int trainedModelsCurrentAllocations,
-        int trainedModelsFailedAllocations
+        int trainedModelsFailedAllocations,
+        int deploymentsWithFixedAllocations,
+        int deploymentsWithDisabledAdaptiveAllocations
     ) {
-        static final TrainedModelAllocationCounts EMPTY = new TrainedModelAllocationCounts(0, 0, 0);
+        static final TrainedModelAllocationCounts EMPTY = new TrainedModelAllocationCounts(0, 0, 0, 0, 0);
     }
 }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/adaptiveallocations/AdaptiveAllocationsScaler.java

@@ -239,4 +239,12 @@ public Double getLastMeasuredInferenceTime() {
     public Long getLastMeasuredQueueSize() {
         return lastMeasuredQueueSize;
     }
+
+    public Integer getMinNumberOfAllocations() {
+        return minNumberOfAllocations;
+    }
+
+    public Integer getMaxNumberOfAllocations() {
+        return maxNumberOfAllocations;
+    }
 }

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/adaptiveallocations/AdaptiveAllocationsScalerService.java

@@ -105,7 +105,7 @@ void init() {
                     "es.ml.trained_models.adaptive_allocations.actual_number_of_allocations.current",
                     "the actual number of allocations",
                     "",
-                    () -> observeLong(AdaptiveAllocationsScaler::getNumberOfAllocations)
+                    this::observeAllocationCount
                 )
             );
             metrics.add(
@@ -179,6 +179,19 @@ Collection<DoubleWithAttributes> observeDouble(Function<AdaptiveAllocationsScale
             }
             return observations;
         }
+
+        Collection<LongWithAttributes> observeAllocationCount() {
+            return scalers.values().stream().map(scaler -> {
+                var value = scaler.getNumberOfAllocations();
+                var min = scaler.getMinNumberOfAllocations();
+                var scalesToZero = min == null || min == 0;
+
+                return new LongWithAttributes(
+                    value,
+                    Map.ofEntries(Map.entry("deployment_id", scaler.getDeploymentId()), Map.entry("scales_to_zero", scalesToZero))
+                );
+            }).toList();
+        }
     }
 
     /**

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/MlMetricsTests.java

@@ -21,6 +21,7 @@
 import org.elasticsearch.xpack.core.ml.datafeed.DatafeedState;
 import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsConfig;
 import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsState;
+import org.elasticsearch.xpack.core.ml.inference.assignment.AdaptiveAllocationsSettings;
 import org.elasticsearch.xpack.core.ml.inference.assignment.RoutingInfo;
 import org.elasticsearch.xpack.core.ml.inference.assignment.RoutingState;
 import org.elasticsearch.xpack.core.ml.inference.assignment.TrainedModelAssignment;
@@ -146,11 +147,27 @@ public void testFindTrainedModelAllocationCounts() {
             TrainedModelAssignment.Builder.empty(mock(StartTrainedModelDeploymentAction.TaskParams.class), null)
                 .addRoutingEntry("node2", new RoutingInfo(0, 1, RoutingState.STARTING, ""))
         );
+        metadataBuilder.addNewAssignment(
+            "model4",
+            TrainedModelAssignment.Builder.empty(
+                mock(StartTrainedModelDeploymentAction.TaskParams.class),
+                new AdaptiveAllocationsSettings(true, 0, 1)
+            ).addRoutingEntry("node1", new RoutingInfo(0, 0, RoutingState.STARTING, ""))
+        );
+        metadataBuilder.addNewAssignment(
+            "model5",
+            TrainedModelAssignment.Builder.empty(
+                mock(StartTrainedModelDeploymentAction.TaskParams.class),
+                new AdaptiveAllocationsSettings(false, 1, 1)
+            ).addRoutingEntry("node1", new RoutingInfo(1, 1, RoutingState.STARTING, ""))
+        );
 
         MlMetrics.TrainedModelAllocationCounts counts = MlMetrics.findTrainedModelAllocationCounts(metadataBuilder.build());
-        assertThat(counts.trainedModelsTargetAllocations(), is(5));
-        assertThat(counts.trainedModelsCurrentAllocations(), is(3));
+        assertThat(counts.trainedModelsTargetAllocations(), is(6));
+        assertThat(counts.trainedModelsCurrentAllocations(), is(4));
         assertThat(counts.trainedModelsFailedAllocations(), is(1));
+        assertThat(counts.deploymentsWithFixedAllocations(), is(3));
+        assertThat(counts.deploymentsWithDisabledAdaptiveAllocations(), is(1));
     }
 
     public void testFindNativeMemoryFree() {