Browse Source

[ML] Audit rebalance message after deployment allocations are increased (#90844)

This commit adds an audit message that a rebalance was performed after
a deployment is updated to have more allocations than before.
Dimitris Athanasiou 3 years ago
parent
commit
1d4c96c9ae

+ 1 - 0
x-pack/plugin/ml/qa/native-multi-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/PyTorchModelIT.java

@@ -825,6 +825,7 @@ public class PyTorchModelIT extends PyTorchModelRestTestCase {
         updateDeployment(modelId, 2);
 
         assertBusy(() -> assertAllocationCount(modelId, 2));
+        assertSystemNotificationsContain("Rebalanced trained model allocations because [model deployment updated]");
     }
 
     public void testUpdateDeployment_GivenAllocationsAreIncreasedOverResources_AndScalingIsPossible() throws Exception {

+ 11 - 1
x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentClusterService.java

@@ -621,7 +621,17 @@ public class TrainedModelAssignmentClusterService implements ClusterStateListene
                 @Override
                 public void clusterStateProcessed(ClusterState oldState, ClusterState newState) {
                     if (isUpdated) {
-                        listener.onResponse(TrainedModelAssignmentMetadata.fromState(newState).getModelAssignment(modelId));
+                        TrainedModelAssignment updatedAssignment = TrainedModelAssignmentMetadata.fromState(newState)
+                            .getModelAssignment(modelId);
+                        if (updatedAssignment.totalTargetAllocations() > existingAssignment.totalTargetAllocations()) {
+                            threadPool.executor(MachineLearning.UTILITY_THREAD_POOL_NAME)
+                                .execute(
+                                    () -> systemAuditor.info(
+                                        Messages.getMessage(Messages.INFERENCE_DEPLOYMENT_REBALANCED, "model deployment updated")
+                                    )
+                                );
+                        }
+                        listener.onResponse(updatedAssignment);
                     }
                 }
             }),