Revert "[ML] Dynamically get of num allocations (#114636)" (#114861)

This reverts commit 8040fbb.

Author: mark-vieira

docs/changelog/114636.yaml

@@ -210,8 +210,4 @@ default List<DefaultConfigId> defaultConfigIds() {
     default void defaultConfigs(ActionListener<List<Model>> defaultsListener) {
         defaultsListener.onResponse(List.of());
     }
-
-    default void updateModelsWithDynamicFields(List<Model> model, ActionListener<List<Model>> listener) {
-        listener.onResponse(model);
-    }
 }

server/src/main/java/org/elasticsearch/inference/InferenceService.java

@@ -9,13 +9,13 @@
 
 import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.ActionRunnable;
 import org.elasticsearch.action.support.ActionFilters;
-import org.elasticsearch.action.support.GroupedActionListener;
 import org.elasticsearch.action.support.HandledTransportAction;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.inference.InferenceServiceRegistry;
-import org.elasticsearch.inference.Model;
+import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.inference.UnparsedModel;
 import org.elasticsearch.injection.guice.Inject;
@@ -29,11 +29,8 @@
 import org.elasticsearch.xpack.inference.registry.ModelRegistry;
 
 import java.util.ArrayList;
-import java.util.Comparator;
-import java.util.HashMap;
 import java.util.List;
 import java.util.concurrent.Executor;
-import java.util.stream.Collectors;
 
 public class TransportGetInferenceModelAction extends HandledTransportAction<
     GetInferenceModelAction.Request,
@@ -99,69 +96,38 @@ private void getSingleModel(
 
             var model = service.get()
                 .parsePersistedConfig(unparsedModel.inferenceEntityId(), unparsedModel.taskType(), unparsedModel.settings());
-
-            service.get()
-                .updateModelsWithDynamicFields(
-                    List.of(model),
-                    delegate.delegateFailureAndWrap(
-                        (l2, updatedModels) -> l2.onResponse(
-                            new GetInferenceModelAction.Response(
-                                updatedModels.stream().map(Model::getConfigurations).collect(Collectors.toList())
-                            )
-                        )
-                    )
-                );
+            delegate.onResponse(new GetInferenceModelAction.Response(List.of(model.getConfigurations())));
         }));
     }
 
     private void getAllModels(ActionListener<GetInferenceModelAction.Response> listener) {
-        modelRegistry.getAllModels(listener.delegateFailureAndWrap((l, models) -> executor.execute(() -> parseModels(models, listener))));
+        modelRegistry.getAllModels(
+            listener.delegateFailureAndWrap((l, models) -> executor.execute(ActionRunnable.supply(l, () -> parseModels(models))))
+        );
     }
 
     private void getModelsByTaskType(TaskType taskType, ActionListener<GetInferenceModelAction.Response> listener) {
         modelRegistry.getModelsByTaskType(
             taskType,
-            listener.delegateFailureAndWrap((l, models) -> executor.execute(() -> parseModels(models, listener)))
+            listener.delegateFailureAndWrap((l, models) -> executor.execute(ActionRunnable.supply(l, () -> parseModels(models))))
         );
     }
 
-    private void parseModels(List<UnparsedModel> unparsedModels, ActionListener<GetInferenceModelAction.Response> listener) {
-        var parsedModelsByService = new HashMap<String, List<Model>>();
-        try {
-            for (var unparsedModel : unparsedModels) {
-                var service = serviceRegistry.getService(unparsedModel.service());
-                if (service.isEmpty()) {
-                    throw serviceNotFoundException(unparsedModel.service(), unparsedModel.inferenceEntityId());
-                }
-                var list = parsedModelsByService.computeIfAbsent(service.get().name(), s -> new ArrayList<>());
-                list.add(
-                    service.get()
-                        .parsePersistedConfig(unparsedModel.inferenceEntityId(), unparsedModel.taskType(), unparsedModel.settings())
-                );
-            }
-
-            var groupedListener = new GroupedActionListener<List<Model>>(
-                parsedModelsByService.entrySet().size(),
-                listener.delegateFailureAndWrap((delegate, listOfListOfModels) -> {
-                    var modifiable = new ArrayList<Model>();
-                    for (var l : listOfListOfModels) {
-                        modifiable.addAll(l);
-                    }
-                    modifiable.sort(Comparator.comparing(Model::getInferenceEntityId));
-                    delegate.onResponse(
-                        new GetInferenceModelAction.Response(modifiable.stream().map(Model::getConfigurations).collect(Collectors.toList()))
-                    );
-                })
-            );
+    private GetInferenceModelAction.Response parseModels(List<UnparsedModel> unparsedModels) {
+        var parsedModels = new ArrayList<ModelConfigurations>();
 
-            for (var entry : parsedModelsByService.entrySet()) {
-                serviceRegistry.getService(entry.getKey())
-                    .get()  // must be non-null to get this far
-                    .updateModelsWithDynamicFields(entry.getValue(), groupedListener);
+        for (var unparsedModel : unparsedModels) {
+            var service = serviceRegistry.getService(unparsedModel.service());
+            if (service.isEmpty()) {
+                throw serviceNotFoundException(unparsedModel.service(), unparsedModel.inferenceEntityId());
             }
-        } catch (Exception e) {
-            listener.onFailure(e);
+            parsedModels.add(
+                service.get()
+                    .parsePersistedConfig(unparsedModel.inferenceEntityId(), unparsedModel.taskType(), unparsedModel.settings())
+                    .getConfigurations()
+            );
         }
+        return new GetInferenceModelAction.Response(parsedModels);
     }
 
     private ElasticsearchStatusException serviceNotFoundException(String service, String inferenceId) {

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

@@ -21,7 +21,7 @@
 
 public abstract class ElasticsearchInternalModel extends Model {
 
-    protected ElasticsearchInternalServiceSettings internalServiceSettings;
+    protected final ElasticsearchInternalServiceSettings internalServiceSettings;
 
     public ElasticsearchInternalModel(
         String inferenceEntityId,
@@ -91,15 +91,6 @@ public ElasticsearchInternalServiceSettings getServiceSettings() {
         return (ElasticsearchInternalServiceSettings) super.getServiceSettings();
     }
 
-    public void updateNumAllocation(Integer numAllocations) {
-        this.internalServiceSettings = new ElasticsearchInternalServiceSettings(
-            numAllocations,
-            this.internalServiceSettings.getNumThreads(),
-            this.internalServiceSettings.modelId(),
-            this.internalServiceSettings.getAdaptiveAllocationsSettings()
-        );
-    }
-
     @Override
     public String toString() {
         return Strings.toString(this.getConfigurations());

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

@@ -32,7 +32,6 @@
 import org.elasticsearch.xpack.core.inference.results.InferenceTextEmbeddingFloatResults;
 import org.elasticsearch.xpack.core.inference.results.RankedDocsResults;
 import org.elasticsearch.xpack.core.inference.results.SparseEmbeddingResults;
-import org.elasticsearch.xpack.core.ml.action.GetDeploymentStatsAction;
 import org.elasticsearch.xpack.core.ml.action.GetTrainedModelsAction;
 import org.elasticsearch.xpack.core.ml.action.GetTrainedModelsStatsAction;
 import org.elasticsearch.xpack.core.ml.action.InferModelAction;
@@ -57,7 +56,6 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.EnumSet;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -788,47 +786,11 @@ public List<DefaultConfigId> defaultConfigIds() {
         );
     }
 
+    /**
+     * Default configurations that can be out of the box without creating an endpoint first.
+     * @param defaultsListener Config listener
+     */
     @Override
-    public void updateModelsWithDynamicFields(List<Model> models, ActionListener<List<Model>> listener) {
-        var modelsByDeploymentIds = new HashMap<String, ElasticsearchInternalModel>();
-        for (var model : models) {
-            if (model instanceof ElasticsearchInternalModel esModel) {
-                modelsByDeploymentIds.put(esModel.internalServiceSettings.deloymentId(), esModel);
-            } else {
-                listener.onFailure(
-                    new ElasticsearchStatusException(
-                        "Cannot update model [{}] as it is not an Elasticsearch service model",
-                        RestStatus.INTERNAL_SERVER_ERROR,
-                        model.getInferenceEntityId()
-                    )
-                );
-                return;
-            }
-        }
-
-        if (modelsByDeploymentIds.isEmpty()) {
-            listener.onResponse(models);
-            return;
-        }
-
-        String deploymentIds = String.join(",", modelsByDeploymentIds.keySet());
-        client.execute(
-            GetDeploymentStatsAction.INSTANCE,
-            new GetDeploymentStatsAction.Request(deploymentIds),
-            ActionListener.wrap(stats -> {
-                for (var deploymentStats : stats.getStats().results()) {
-                    var model = modelsByDeploymentIds.get(deploymentStats.getDeploymentId());
-                    model.updateNumAllocation(deploymentStats.getNumberOfAllocations());
-                }
-                listener.onResponse(new ArrayList<>(modelsByDeploymentIds.values()));
-            }, e -> {
-                logger.warn("Get deployment stats failed, cannot update the endpoint's number of allocations", e);
-                // continue with the original response
-                listener.onResponse(models);
-            })
-        );
-    }
-
     public void defaultConfigs(ActionListener<List<Model>> defaultsListener) {
         preferredModelVariantFn.accept(defaultsListener.delegateFailureAndWrap((delegate, preferredModelVariant) -> {
             if (PreferredModelVariant.LINUX_X86_OPTIMIZED.equals(preferredModelVariant)) {

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

@@ -194,10 +194,6 @@ public String modelId() {
         return modelId;
     }
 
-    public String deloymentId() {
-        return modelId;
-    }
-
     public Integer getNumAllocations() {
         return numAllocations;
     }