[ML] Block trained model updates from inference (elastic#130940)

When the Trained Model has been deployed through the Inference Endpoint
API, it can only be updated using the Inference Endpoint API.

When the Trained Model has been deployed and then attached to an
Inference Endpoint, it can only be updated using the Trained Model API.

Fix elastic#129999

Co-authored-by: elasticsearchmachine <[email protected]>
Co-authored-by: David Kyle <[email protected]>

Author: 3 people

docs/changelog/130940.yaml

@@ -0,0 +1,6 @@
+pr: 130940
+summary: Block trained model updates from inference
+area: Machine Learning
+type: enhancement
+issues:
+ - 129999

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/job/messages/Messages.java

@@ -286,6 +286,12 @@ public final class Messages {
     public static final String MODEL_ID_DOES_NOT_MATCH_EXISTING_MODEL_IDS_BUT_MUST_FOR_IN_CLUSTER_SERVICE =
         "Requested model ID [{}] does not have a matching trained model and thus cannot be updated.";
     public static final String INFERENCE_ENTITY_NON_EXISTANT_NO_UPDATE = "The inference endpoint [{}] does not exist and cannot be updated";
+    public static final String INFERENCE_REFERENCE_CANNOT_UPDATE_ANOTHER_ENDPOINT =
+        "Cannot update inference endpoint [{}] for model deployment [{}] as it was created by another inference endpoint. "
+            + "The model can only be updated using inference endpoint id [{}].";
+    public static final String INFERENCE_CAN_ONLY_UPDATE_MODELS_IT_CREATED =
+        "Cannot update inference endpoint [{}] using model deployment [{}]. "
+            + "The model deployment must be updated through the trained models API.";
 
     private Messages() {}
 

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/CreateFromDeploymentIT.java

@@ -51,24 +51,6 @@ public void testAttachToDeployment() throws IOException {
         var results = infer(inferenceId, List.of("washing machine"));
         assertNotNull(results.get("sparse_embedding"));
 
-        var updatedNumAllocations = randomIntBetween(1, 10);
-        var updatedEndpointConfig = updateEndpoint(inferenceId, updatedEndpointConfig(updatedNumAllocations), TaskType.SPARSE_EMBEDDING);
-        assertThat(
-            updatedEndpointConfig.get("service_settings"),
-            is(
-                Map.of(
-                    "num_allocations",
-                    updatedNumAllocations,
-                    "num_threads",
-                    1,
-                    "model_id",
-                    "attach_to_deployment",
-                    "deployment_id",
-                    "existing_deployment"
-                )
-            )
-        );
-
         deleteModel(inferenceId);
         // assert deployment not stopped
         var stats = (List<Map<String, Object>>) getTrainedModelStats(modelId).get("trained_model_stats");
@@ -128,24 +110,6 @@ public void testAttachWithModelId() throws IOException {
         var results = infer(inferenceId, List.of("washing machine"));
         assertNotNull(results.get("sparse_embedding"));
 
-        var updatedNumAllocations = randomIntBetween(1, 10);
-        var updatedEndpointConfig = updateEndpoint(inferenceId, updatedEndpointConfig(updatedNumAllocations), TaskType.SPARSE_EMBEDDING);
-        assertThat(
-            updatedEndpointConfig.get("service_settings"),
-            is(
-                Map.of(
-                    "num_allocations",
-                    updatedNumAllocations,
-                    "num_threads",
-                    1,
-                    "model_id",
-                    "attach_with_model_id",
-                    "deployment_id",
-                    "existing_deployment_with_model_id"
-                )
-            )
-        );
-
         forceStopMlNodeDeployment(deploymentId);
     }
 
@@ -180,6 +144,30 @@ public void testDeploymentDoesNotExist() {
         assertThat(e.getMessage(), containsString("Cannot find deployment [missing_deployment]"));
     }
 
+    public void testCreateInferenceUsingSameDeploymentId() throws IOException {
+        var modelId = "conflicting_ids";
+        var deploymentId = modelId;
+        var inferenceId = modelId;
+
+        CustomElandModelIT.createMlNodeTextExpansionModel(modelId, client());
+        var response = startMlNodeDeploymemnt(modelId, deploymentId);
+        assertStatusOkOrCreated(response);
+
+        var responseException = assertThrows(
+            ResponseException.class,
+            () -> putModel(inferenceId, endpointConfig(deploymentId), TaskType.SPARSE_EMBEDDING)
+        );
+        assertThat(
+            responseException.getMessage(),
+            containsString(
+                "Inference endpoint IDs must be unique. "
+                    + "Requested inference endpoint ID [conflicting_ids] matches existing trained model ID(s) but must not."
+            )
+        );
+
+        forceStopMlNodeDeployment(deploymentId);
+    }
+
     public void testNumAllocationsIsUpdated() throws IOException {
         var modelId = "update_num_allocations";
         var deploymentId = modelId;
@@ -208,7 +196,16 @@ public void testNumAllocationsIsUpdated() throws IOException {
             )
         );
 
-        assertStatusOkOrCreated(updateMlNodeDeploymemnt(deploymentId, 2));
+        var responseException = assertThrows(ResponseException.class, () -> updateInference(inferenceId, TaskType.SPARSE_EMBEDDING, 2));
+        assertThat(
+            responseException.getMessage(),
+            containsString(
+                "Cannot update inference endpoint [test_num_allocations_updated] using model deployment [update_num_allocations]. "
+                    + "The model deployment must be updated through the trained models API."
+            )
+        );
+
+        updateMlNodeDeploymemnt(deploymentId, 2);
 
         var updatedServiceSettings = getModel(inferenceId).get("service_settings");
         assertThat(
@@ -227,6 +224,92 @@ public void testNumAllocationsIsUpdated() throws IOException {
                 )
             )
         );
+
+        forceStopMlNodeDeployment(deploymentId);
+    }
+
+    public void testUpdateWhenInferenceEndpointCreatesDeployment() throws IOException {
+        var modelId = "update_num_allocations_from_created_endpoint";
+        var inferenceId = "test_created_endpoint_from_model";
+        var deploymentId = inferenceId;
+
+        CustomElandModelIT.createMlNodeTextExpansionModel(modelId, client());
+
+        var putModel = putModel(inferenceId, Strings.format("""
+            {
+              "service": "elasticsearch",
+              "service_settings": {
+                "num_allocations": %s,
+                "num_threads": %s,
+                "model_id": "%s"
+              }
+            }
+            """, 1, 1, modelId), TaskType.SPARSE_EMBEDDING);
+        var serviceSettings = putModel.get("service_settings");
+        assertThat(putModel.toString(), serviceSettings, is(Map.of("num_allocations", 1, "num_threads", 1, "model_id", modelId)));
+
+        updateInference(inferenceId, TaskType.SPARSE_EMBEDDING, 2);
+
+        var responseException = assertThrows(ResponseException.class, () -> updateMlNodeDeploymemnt(deploymentId, 2));
+        assertThat(
+            responseException.getMessage(),
+            containsString(
+                "Cannot update deployment [test_created_endpoint_from_model] as it was created by inference endpoint "
+                    + "[test_created_endpoint_from_model]. This model deployment must be updated through the inference API."
+            )
+        );
+
+        var updatedServiceSettings = getModel(inferenceId).get("service_settings");
+        assertThat(
+            updatedServiceSettings.toString(),
+            updatedServiceSettings,
+            is(Map.of("num_allocations", 2, "num_threads", 1, "model_id", modelId))
+        );
+
+        forceStopMlNodeDeployment(deploymentId);
+    }
+
+    public void testCannotUpdateAnotherInferenceEndpointsCreatedDeployment() throws IOException {
+        var modelId = "model_deployment_for_endpoint";
+        var inferenceId = "first_endpoint_for_model_deployment";
+        var deploymentId = inferenceId;
+
+        CustomElandModelIT.createMlNodeTextExpansionModel(modelId, client());
+
+        putModel(inferenceId, Strings.format("""
+            {
+              "service": "elasticsearch",
+              "service_settings": {
+                "num_allocations": %s,
+                "num_threads": %s,
+                "model_id": "%s"
+              }
+            }
+            """, 1, 1, modelId), TaskType.SPARSE_EMBEDDING);
+
+        var secondInferenceId = "second_endpoint_for_model_deployment";
+        var putModel = putModel(secondInferenceId, endpointConfig(deploymentId), TaskType.SPARSE_EMBEDDING);
+        var serviceSettings = putModel.get("service_settings");
+        assertThat(
+            putModel.toString(),
+            serviceSettings,
+            is(Map.of("num_allocations", 1, "num_threads", 1, "model_id", modelId, "deployment_id", deploymentId))
+        );
+
+        var responseException = assertThrows(
+            ResponseException.class,
+            () -> updateInference(secondInferenceId, TaskType.SPARSE_EMBEDDING, 2)
+        );
+        assertThat(
+            responseException.getMessage(),
+            containsString(
+                "Cannot update inference endpoint [second_endpoint_for_model_deployment] for model deployment "
+                    + "[first_endpoint_for_model_deployment] as it was created by another inference endpoint. "
+                    + "The model can only be updated using inference endpoint id [first_endpoint_for_model_deployment]."
+            )
+        );
+
+        forceStopMlNodeDeployment(deploymentId);
     }
 
     public void testStoppingDeploymentAttachedToInferenceEndpoint() throws IOException {
@@ -300,6 +383,22 @@ private Response startMlNodeDeploymemnt(String modelId, String deploymentId) thr
         return client().performRequest(request);
     }
 
+    private Response updateInference(String deploymentId, TaskType taskType, int numAllocations) throws IOException {
+        String endPoint = Strings.format("/_inference/%s/%s/_update", taskType, deploymentId);
+
+        var body = Strings.format("""
+            {
+              "service_settings": {
+                "num_allocations": %d
+              }
+            }
+            """, numAllocations);
+
+        Request request = new Request("PUT", endPoint);
+        request.setJsonEntity(body);
+        return client().performRequest(request);
+    }
+
     private Response updateMlNodeDeploymemnt(String deploymentId, int numAllocations) throws IOException {
         String endPoint = "/_ml/trained_models/" + deploymentId + "/deployment/_update";
 
@@ -314,6 +413,16 @@ private Response updateMlNodeDeploymemnt(String deploymentId, int numAllocations
         return client().performRequest(request);
     }
 
+    private Map<String, Object> updateMlNodeDeploymemnt(String deploymentId, String body) throws IOException {
+        String endPoint = "/_ml/trained_models/" + deploymentId + "/deployment/_update";
+
+        Request request = new Request("POST", endPoint);
+        request.setJsonEntity(body);
+        var response = client().performRequest(request);
+        assertStatusOkOrCreated(response);
+        return entityAsMap(response);
+    }
+
     protected void stopMlNodeDeployment(String deploymentId) throws IOException {
         String endpoint = "/_ml/trained_models/" + deploymentId + "/deployment/_stop";
         Request request = new Request("POST", endpoint);

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

@@ -58,7 +58,6 @@
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -224,13 +223,10 @@ private Model combineExistingModelWithNewSettings(
         if (settingsToUpdate.serviceSettings() != null && existingSecretSettings != null) {
             newSecretSettings = existingSecretSettings.newSecretSettings(settingsToUpdate.serviceSettings());
         }
-        if (settingsToUpdate.serviceSettings() != null && settingsToUpdate.serviceSettings().containsKey(NUM_ALLOCATIONS)) {
-            // In cluster services can only have their num_allocations updated, so this is a special case
+        if (settingsToUpdate.serviceSettings() != null) {
+            // In cluster services can have their deployment settings updated, so this is a special case
             if (newServiceSettings instanceof ElasticsearchInternalServiceSettings elasticServiceSettings) {
-                newServiceSettings = new ElasticsearchInternalServiceSettings(
-                    elasticServiceSettings,
-                    (Integer) settingsToUpdate.serviceSettings().get(NUM_ALLOCATIONS)
-                );
+                newServiceSettings = elasticServiceSettings.updateServiceSettings(settingsToUpdate.serviceSettings());
             }
         }
         if (settingsToUpdate.taskSettings() != null && existingTaskSettings != null) {
@@ -257,26 +253,59 @@ private void updateInClusterEndpoint(
         Model newModel,
         Model existingParsedModel,
         ActionListener<Boolean> listener
-    ) throws IOException {
+    ) {
         // The model we are trying to update must have a trained model associated with it if it is an in-cluster deployment
         var deploymentId = getDeploymentIdForInClusterEndpoint(existingParsedModel);
-        throwIfTrainedModelDoesntExist(request.getInferenceEntityId(), deploymentId);
+        var inferenceEntityId = request.getInferenceEntityId();
+        throwIfTrainedModelDoesntExist(inferenceEntityId, deploymentId);
 
-        Map<String, Object> serviceSettings = request.getContentAsSettings().serviceSettings();
-        if (serviceSettings != null && serviceSettings.get(NUM_ALLOCATIONS) instanceof Integer numAllocations) {
+        if (inferenceEntityId.equals(deploymentId) == false) {
+            modelRegistry.getModel(deploymentId, ActionListener.wrap(unparsedModel -> {
+                // if this deployment was created by another inference endpoint, then it must be updated using that inference endpoint
+                listener.onFailure(
+                    new ElasticsearchStatusException(
+                        Messages.INFERENCE_REFERENCE_CANNOT_UPDATE_ANOTHER_ENDPOINT,
+                        RestStatus.CONFLICT,
+                        inferenceEntityId,
+                        deploymentId,
+                        unparsedModel.inferenceEntityId()
+                    )
+                );
+            }, e -> {
+                if (e instanceof ResourceNotFoundException) {
+                    // if this deployment was created by the trained models API, then it must be updated by the trained models API
+                    listener.onFailure(
+                        new ElasticsearchStatusException(
+                            Messages.INFERENCE_CAN_ONLY_UPDATE_MODELS_IT_CREATED,
+                            RestStatus.CONFLICT,
+                            inferenceEntityId,
+                            deploymentId
+                        )
+                    );
+                    return;
+                }
+                listener.onFailure(e);
+            }));
+            return;
+        }
+
+        if (newModel.getServiceSettings() instanceof ElasticsearchInternalServiceSettings elasticServiceSettings) {
 
-            UpdateTrainedModelDeploymentAction.Request updateRequest = new UpdateTrainedModelDeploymentAction.Request(deploymentId);
-            updateRequest.setNumberOfAllocations(numAllocations);
+            var updateRequest = new UpdateTrainedModelDeploymentAction.Request(deploymentId);
+            updateRequest.setNumberOfAllocations(elasticServiceSettings.getNumAllocations());
+            updateRequest.setAdaptiveAllocationsSettings(elasticServiceSettings.getAdaptiveAllocationsSettings());
+            updateRequest.setIsInternal(true);
 
             var delegate = listener.<CreateTrainedModelAssignmentAction.Response>delegateFailure((l2, response) -> {
                 modelRegistry.updateModelTransaction(newModel, existingParsedModel, l2);
             });
 
             logger.info(
-                "Updating trained model deployment [{}] for inference entity [{}] with [{}] num_allocations",
+                "Updating trained model deployment [{}] for inference entity [{}] with [{}] num_allocations and adaptive allocations [{}]",
                 deploymentId,
                 request.getInferenceEntityId(),
-                numAllocations
+                elasticServiceSettings.getNumAllocations(),
+                elasticServiceSettings.getAdaptiveAllocationsSettings()
             );
             client.execute(UpdateTrainedModelDeploymentAction.INSTANCE, updateRequest, delegate);
 
@@ -317,7 +346,6 @@ private void throwIfTrainedModelDoesntExist(String inferenceEntityId, String dep
             throw ExceptionsHelper.entityNotFoundException(
                 Messages.MODEL_ID_DOES_NOT_MATCH_EXISTING_MODEL_IDS_BUT_MUST_FOR_IN_CLUSTER_SERVICE,
                 inferenceEntityId
-
             );
         }
     }