[ML] updating node memory load for new allocation service (#76046)

Trained model deployment memory usage is no longer determinable via persistent tasks.

The new way is to look into the trained model allocation metadata.

This PR updates this and removes some unused code.

relates: #75778

Author: benwtrent

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

@@ -15,8 +15,6 @@
 import org.elasticsearch.xpack.core.ml.datafeed.DatafeedState;
 import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsState;
 import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsTaskState;
-import org.elasticsearch.xpack.core.ml.inference.deployment.TrainedModelDeploymentState;
-import org.elasticsearch.xpack.core.ml.inference.deployment.TrainedModelDeploymentTaskState;
 import org.elasticsearch.xpack.core.ml.job.config.JobState;
 import org.elasticsearch.xpack.core.ml.job.config.JobTaskState;
 import org.elasticsearch.xpack.core.ml.job.snapshot.upgrade.SnapshotUpgradeState;
@@ -30,11 +28,13 @@
 
 public final class MlTasks {
 
+    public static final String TRAINED_MODEL_ALLOCATION_TASK_TYPE = "trained_model_allocation";
+    public static final String TRAINED_MODEL_ALLOCATION_TASK_NAME_PREFIX = "xpack/ml/allocation-";
+
     public static final String JOB_TASK_NAME = "xpack/ml/job";
     public static final String DATAFEED_TASK_NAME = "xpack/ml/datafeed";
     public static final String DATA_FRAME_ANALYTICS_TASK_NAME = "xpack/ml/data_frame/analytics";
     public static final String JOB_SNAPSHOT_UPGRADE_TASK_NAME = "xpack/ml/job/snapshot/upgrade";
-    public static final String TRAINED_MODEL_DEPLOYMENT_TASK_NAME = "xpack/ml/trained_model/deployment";
 
     public static final String JOB_TASK_ID_PREFIX = "job-";
     public static final String DATAFEED_TASK_ID_PREFIX = "datafeed-";
@@ -225,31 +225,6 @@ public static DataFrameAnalyticsState getDataFrameAnalyticsState(@Nullable Persi
         return state;
     }
 
-    public static TrainedModelDeploymentState getTrainedModelDeploymentState(PersistentTasksCustomMetadata.PersistentTask<?> task) {
-        if (task == null) {
-            return TrainedModelDeploymentState.STOPPED;
-        }
-        TrainedModelDeploymentTaskState taskState = (TrainedModelDeploymentTaskState) task.getState();
-        if (taskState == null) {
-            return TrainedModelDeploymentState.STARTING;
-        }
-
-        TrainedModelDeploymentState state = taskState.getState();
-        if (taskState.isStatusStale(task)) {
-            if (state == TrainedModelDeploymentState.STOPPING) {
-                // previous executor node failed while the job was stopping - it won't
-                // be restarted on another node, so consider it STOPPED for reassignment purposes
-                return TrainedModelDeploymentState.STOPPED;
-            }
-            if (state != TrainedModelDeploymentState.FAILED) {
-                // we are relocating at the moment
-                // TODO Revisit this in the new allocation framework as there won't necessarily be a concept of relocation.
-                return TrainedModelDeploymentState.STARTING;
-            }
-        }
-        return state;
-    }
-
     /**
      * The job Ids of anomaly detector job tasks.
      * All anomaly detector jobs are returned regardless of the status of the
@@ -435,8 +410,6 @@ public static MemoryTrackedTaskState getMemoryTrackedTaskState(PersistentTasksCu
                 return taskState == null ? SnapshotUpgradeState.LOADING_OLD_STATE : taskState.getState();
             case DATA_FRAME_ANALYTICS_TASK_NAME:
                 return getDataFrameAnalyticsState(task);
-            case TRAINED_MODEL_DEPLOYMENT_TASK_NAME:
-                return getTrainedModelDeploymentState(task);
             default:
                 throw new IllegalStateException("unexpected task type [" + task.getTaskName() + "]");
         }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/StartTrainedModelDeploymentAction.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.action.support.master.MasterNodeRequest;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodeRole;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.xcontent.ConstructingObjectParser;
 import org.elasticsearch.common.xcontent.ParseField;
@@ -23,7 +24,6 @@
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.common.xcontent.ToXContentObject;
 import org.elasticsearch.common.xcontent.XContentBuilder;
-import org.elasticsearch.persistent.PersistentTaskParams;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.xpack.core.ml.MlTasks;
 import org.elasticsearch.xpack.core.ml.inference.TrainedModelConfig;
@@ -122,13 +122,13 @@ public String toString() {
         }
     }
 
-    public static class TaskParams implements PersistentTaskParams, MlTaskParams {
+    public static class TaskParams implements MlTaskParams, Writeable, ToXContentObject {
 
         // TODO add support for other roles? If so, it may have to be an instance method...
         // NOTE, whatever determines allocation should not be dynamically set on the node
         // Otherwise allocation logic might fail
         public static boolean mayAllocateToNode(DiscoveryNode node) {
-            return node.getRoles().contains(DiscoveryNodeRole.ML_ROLE);
+            return node.getRoles().contains(DiscoveryNodeRole.ML_ROLE) && node.getVersion().onOrAfter(VERSION_INTRODUCED);
         }
 
         public static final Version VERSION_INTRODUCED = Version.V_8_0_0;
@@ -187,12 +187,6 @@ public long estimateMemoryUsageBytes() {
             return MEMORY_OVERHEAD.getBytes() + 2 * modelBytes;
         }
 
-        @Override
-        public String getWriteableName() {
-            return MlTasks.TRAINED_MODEL_DEPLOYMENT_TASK_NAME;
-        }
-
-        @Override
         public Version getMinimalSupportedVersion() {
             return VERSION_INTRODUCED;
         }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/allocation/RoutingState.java

@@ -7,9 +7,12 @@
 
 package org.elasticsearch.xpack.core.ml.inference.allocation;
 
+import org.elasticsearch.xpack.core.ml.utils.MemoryTrackedTaskState;
+
+import java.util.Arrays;
 import java.util.Locale;
 
-public enum RoutingState {
+public enum RoutingState implements MemoryTrackedTaskState {
     STARTING,
     STARTED,
     STOPPING,
@@ -20,8 +23,25 @@ public static RoutingState fromString(String value) {
         return valueOf(value.toUpperCase(Locale.ROOT));
     }
 
+    /**
+     * @return {@code true} if state matches none of the given {@code candidates}
+     */
+    public boolean isNoneOf(RoutingState... candidates) {
+        return Arrays.stream(candidates).noneMatch(candidate -> this == candidate);
+    }
+
     @Override
     public String toString() {
         return name().toLowerCase(Locale.ROOT);
     }
+
+    @Override
+    public boolean consumesMemory() {
+        return isNoneOf(FAILED, STOPPED);
+    }
+
+    @Override
+    public boolean isAllocating() {
+        return this == STARTING;
+    }
 }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/deployment/TrainedModelDeploymentState.java

@@ -16,17 +16,13 @@
 import org.elasticsearch.xpack.core.ml.action.OpenJobAction;
 import org.elasticsearch.xpack.core.ml.action.StartDataFrameAnalyticsAction;
 import org.elasticsearch.xpack.core.ml.action.StartDatafeedAction;
-import org.elasticsearch.xpack.core.ml.action.StartTrainedModelDeploymentAction;
 import org.elasticsearch.xpack.core.ml.datafeed.DatafeedState;
 import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsState;
 import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsTaskState;
-import org.elasticsearch.xpack.core.ml.inference.deployment.TrainedModelDeploymentState;
-import org.elasticsearch.xpack.core.ml.inference.deployment.TrainedModelDeploymentTaskState;
 import org.elasticsearch.xpack.core.ml.job.config.JobState;
 import org.elasticsearch.xpack.core.ml.job.config.JobTaskState;
 
 import java.net.InetAddress;
-import java.util.Arrays;
 
 import static org.hamcrest.Matchers.contains;
 import static org.hamcrest.Matchers.containsInAnyOrder;
@@ -308,41 +304,6 @@ public void testGetDataFrameAnalyticsState_GivenStaleTaskWithFailedState() {
         assertThat(state, equalTo(DataFrameAnalyticsState.FAILED));
     }
 
-    public void testGetTrainedModelDeploymentState_GivenNull() {
-        assertThat(MlTasks.getTrainedModelDeploymentState(null), equalTo(TrainedModelDeploymentState.STOPPED));
-    }
-
-    public void testGetTrainedModelDeploymentState_GivenTaskStateIsNull() {
-        PersistentTasksCustomMetadata.PersistentTask<?> task = createTrainedModelTask(null, false);
-        assertThat(MlTasks.getTrainedModelDeploymentState(task), equalTo(TrainedModelDeploymentState.STARTING));
-    }
-
-    public void testGetTrainedModelDeploymentState_GivenTaskStateIsNotNullAndNotStale() {
-        TrainedModelDeploymentState state = randomFrom(TrainedModelDeploymentState.values());
-        PersistentTasksCustomMetadata.PersistentTask<?> task = createTrainedModelTask(state, false);
-        assertThat(MlTasks.getTrainedModelDeploymentState(task), equalTo(state));
-    }
-
-    public void testGetTrainedModelDeploymentState_GivenTaskStateIsStaleAndStopping() {
-        PersistentTasksCustomMetadata.PersistentTask<?> task = createTrainedModelTask(TrainedModelDeploymentState.STOPPING, true);
-        assertThat(MlTasks.getTrainedModelDeploymentState(task), equalTo(TrainedModelDeploymentState.STOPPED));
-    }
-
-    public void testGetTrainedModelDeploymentState_GivenTaskStateIsStaleAndFailed() {
-        PersistentTasksCustomMetadata.PersistentTask<?> task = createTrainedModelTask(TrainedModelDeploymentState.FAILED, true);
-        assertThat(MlTasks.getTrainedModelDeploymentState(task), equalTo(TrainedModelDeploymentState.FAILED));
-    }
-
-    public void testGetTrainedModelDeploymentState_GivenTaskStateIsStaleAndNotFailedNorStopping() {
-        TrainedModelDeploymentState state = randomFrom(
-            Arrays.stream(TrainedModelDeploymentState.values())
-                .filter(s -> s != TrainedModelDeploymentState.FAILED && s != TrainedModelDeploymentState.STOPPING)
-                .toArray(TrainedModelDeploymentState[]::new)
-        );
-        PersistentTasksCustomMetadata.PersistentTask<?> task = createTrainedModelTask(state, true);
-        assertThat(MlTasks.getTrainedModelDeploymentState(task), equalTo(TrainedModelDeploymentState.STARTING));
-    }
-
     private static PersistentTasksCustomMetadata.PersistentTask<?> createDataFrameAnalyticsTask(String jobId, String nodeId,
                                                                                                 DataFrameAnalyticsState state,
                                                                                                 boolean isStale) {
@@ -358,18 +319,4 @@ private static PersistentTasksCustomMetadata.PersistentTask<?> createDataFrameAn
         return tasks.getTask(MlTasks.dataFrameAnalyticsTaskId(jobId));
     }
 
-    private static PersistentTasksCustomMetadata.PersistentTask<?> createTrainedModelTask(TrainedModelDeploymentState state,
-                                                                                          boolean isStale) {
-        String id = randomAlphaOfLength(10);
-        PersistentTasksCustomMetadata.Builder builder = PersistentTasksCustomMetadata.builder();
-        builder.addTask(MlTasks.trainedModelDeploymentTaskId(id), MlTasks.TRAINED_MODEL_DEPLOYMENT_TASK_NAME,
-            new StartTrainedModelDeploymentAction.TaskParams(id, randomAlphaOfLength(10), randomNonNegativeLong()),
-            new PersistentTasksCustomMetadata.Assignment(randomAlphaOfLength(10), "test assignment"));
-        if (state != null) {
-            builder.updateTaskState(MlTasks.trainedModelDeploymentTaskId(id),
-                new TrainedModelDeploymentTaskState(state, builder.getLastAllocationId() - (isStale ? 1 : 0), null));
-        }
-        PersistentTasksCustomMetadata tasks = builder.build();
-        return tasks.getTask(MlTasks.trainedModelDeploymentTaskId(id));
-    }
 }