Fix scale up for model allocations (#115189) (#115225)

elastic · Oct 21, 2024 · a312404 · a312404
1 parent 5cc18bc
commit a312404
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 1 deletion.
diff --git a/.../plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlAutoscalingContext.java b/.../plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/MlAutoscalingContext.java
@@ -177,7 +177,7 @@ public boolean isEmpty() {
         return anomalyDetectionTasks.isEmpty()
             && snapshotUpgradeTasks.isEmpty()
             && dataframeAnalyticsTasks.isEmpty()
-            && modelAssignments.values().stream().allMatch(assignment -> assignment.totalTargetAllocations() == 0);
+            && modelAssignments.values().stream().allMatch(assignment -> assignment.getTaskParams().getNumberOfAllocations() == 0);
     }
 
     public List<String> findPartiallyAllocatedModels() {

diff --git a/...rc/test/java/org/elasticsearch/xpack/ml/autoscaling/MlAutoscalingDeciderServiceTests.java b/...rc/test/java/org/elasticsearch/xpack/ml/autoscaling/MlAutoscalingDeciderServiceTests.java
@@ -54,6 +54,7 @@
 import static org.elasticsearch.xpack.ml.utils.NativeMemoryCalculator.STATIC_JVM_UPPER_THRESHOLD;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.nullValue;
 import static org.mockito.ArgumentMatchers.any;
@@ -331,6 +332,53 @@ public void testScale_GivenModelWithZeroAllocations() {
         assertThat(result.requiredCapacity().node().memory().getBytes(), equalTo(0L));
     }
 
+    public void testScale_GivenTrainedModelAllocationAndNoMlNode() {
+        MlAutoscalingDeciderService service = buildService();
+        service.onMaster();
+
+        ClusterState clusterState = new ClusterState.Builder(new ClusterName("cluster")).metadata(
+            Metadata.builder()
+                .putCustom(
+                    TrainedModelAssignmentMetadata.NAME,
+                    new TrainedModelAssignmentMetadata(
+                        Map.of(
+                            "model",
+                            TrainedModelAssignment.Builder.empty(
+                                new StartTrainedModelDeploymentAction.TaskParams(
+                                    "model",
+                                    "model-deployment",
+                                    400,
+                                    1,
+                                    2,
+                                    100,
+                                    null,
+                                    Priority.NORMAL,
+                                    0L,
+                                    0L
+                                ),
+                                new AdaptiveAllocationsSettings(true, 0, 4)
+                            ).setAssignmentState(AssignmentState.STARTING).build()
+                        )
+                    )
+                )
+                .build()
+        ).build();
+
+        AutoscalingDeciderResult result = service.scale(
+            Settings.EMPTY,
+            new DeciderContext(
+                clusterState,
+                new AutoscalingCapacity(AutoscalingCapacity.AutoscalingResources.ZERO, AutoscalingCapacity.AutoscalingResources.ZERO)
+            )
+        );
+
+        assertThat(result.reason().summary(), containsString("requesting scale up"));
+        assertThat(result.requiredCapacity().total().memory().getBytes(), greaterThan(TEST_JOB_SIZE));
+        assertThat(result.requiredCapacity().total().processors().count(), equalTo(2.0));
+        assertThat(result.requiredCapacity().node().memory().getBytes(), greaterThan(TEST_JOB_SIZE));
+        assertThat(result.requiredCapacity().node().processors().count(), equalTo(2.0));
+    }
+
     private DiscoveryNode buildNode(String id, ByteSizeValue machineMemory, int allocatedProcessors) {
         return DiscoveryNodeUtils.create(
             id,