Browse Source

[ML] Improve trained model allocation task info (#81259)

This commit changes the trained model allocation task info
in two ways:

  - the action is now `xpack/ml/trained_model_allocation` and does
  not contain the model_id as this is consistent with other ES tasks.
  - the description is now `model_id[<model_id>]`. We do not need to
  prefix this with anything and it follows other ES tasks where information
  is added in square brackets following the information description.
Dimitris Athanasiou 3 years ago
parent
commit
89c3915711

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

@@ -15,6 +15,7 @@ import org.elasticsearch.persistent.PersistentTasksCustomMetadata;
 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.TrainedModelConfig;
 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;
@@ -29,7 +30,7 @@ import java.util.stream.Collectors;
 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 TRAINED_MODEL_ALLOCATION_TASK_ACTION = "xpack/ml/trained_model_allocation[n]";
 
     public static final String JOB_TASK_NAME = "xpack/ml/job";
     public static final String DATAFEED_TASK_NAME = "xpack/ml/datafeed";
@@ -40,7 +41,6 @@ public final class MlTasks {
     public static final String DATAFEED_TASK_ID_PREFIX = "datafeed-";
     public static final String DATA_FRAME_ANALYTICS_TASK_ID_PREFIX = "data_frame_analytics-";
     public static final String JOB_SNAPSHOT_UPGRADE_TASK_ID_PREFIX = "job-snapshot-upgrade-";
-    public static final String TRAINED_MODEL_DEPLOYMENT_TASK_ID_PREFIX = "trained_model_deployment-";
 
     public static final PersistentTasksCustomMetadata.Assignment AWAITING_UPGRADE = new PersistentTasksCustomMetadata.Assignment(
         null,
@@ -99,12 +99,8 @@ public final class MlTasks {
         return taskId.substring(DATA_FRAME_ANALYTICS_TASK_ID_PREFIX.length());
     }
 
-    public static String trainedModelDeploymentTaskId(String deploymentId) {
-        return TRAINED_MODEL_DEPLOYMENT_TASK_ID_PREFIX + deploymentId;
-    }
-
-    public static String trainedModelDeploymentId(String taskId) {
-        return taskId.substring(TRAINED_MODEL_DEPLOYMENT_TASK_ID_PREFIX.length());
+    public static String trainedModelAllocationTaskDescription(String modelId) {
+        return TrainedModelConfig.MODEL_ID.getPreferredName() + "[" + modelId + "]";
     }
 
     @Nullable

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

@@ -36,7 +36,7 @@ import java.io.IOException;
 import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 
-import static org.elasticsearch.xpack.core.ml.MlTasks.trainedModelDeploymentTaskId;
+import static org.elasticsearch.xpack.core.ml.MlTasks.trainedModelAllocationTaskDescription;
 
 public class StartTrainedModelDeploymentAction extends ActionType<CreateTrainedModelAllocationAction.Response> {
 
@@ -382,7 +382,7 @@ public class StartTrainedModelDeploymentAction extends ActionType<CreateTrainedM
                 if (Strings.isAllOrWildcard(expectedId)) {
                     return true;
                 }
-                String expectedDescription = trainedModelDeploymentTaskId(expectedId);
+                String expectedDescription = trainedModelAllocationTaskDescription(expectedId);
                 return expectedDescription.equals(task.getDescription());
             }
             return false;

+ 2 - 2
x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/allocation/TrainedModelAllocationNodeService.java

@@ -51,7 +51,7 @@ import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedDeque;
 
-import static org.elasticsearch.xpack.core.ml.MlTasks.TRAINED_MODEL_ALLOCATION_TASK_NAME_PREFIX;
+import static org.elasticsearch.xpack.core.ml.MlTasks.TRAINED_MODEL_ALLOCATION_TASK_ACTION;
 import static org.elasticsearch.xpack.core.ml.MlTasks.TRAINED_MODEL_ALLOCATION_TASK_TYPE;
 import static org.elasticsearch.xpack.ml.MachineLearning.ML_PYTORCH_MODEL_INFERENCE_FEATURE;
 
@@ -356,7 +356,7 @@ public class TrainedModelAllocationNodeService implements ClusterStateListener {
         );
         TrainedModelDeploymentTask task = (TrainedModelDeploymentTask) taskManager.register(
             TRAINED_MODEL_ALLOCATION_TASK_TYPE,
-            TRAINED_MODEL_ALLOCATION_TASK_NAME_PREFIX + taskParams.getModelId(),
+            TRAINED_MODEL_ALLOCATION_TASK_ACTION,
             taskAwareRequest(taskParams)
         );
         // threadsafe check to verify we are not loading/loaded the model

+ 1 - 1
x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/deployment/TrainedModelDeploymentTask.java

@@ -51,7 +51,7 @@ public class TrainedModelDeploymentTask extends CancellableTask implements Start
         XPackLicenseState licenseState,
         LicensedFeature.Persistent licensedFeature
     ) {
-        super(id, type, action, MlTasks.trainedModelDeploymentTaskId(taskParams.getModelId()), parentTask, headers);
+        super(id, type, action, MlTasks.trainedModelAllocationTaskDescription(taskParams.getModelId()), parentTask, headers);
         this.params = taskParams;
         this.trainedModelAllocationNodeService = ExceptionsHelper.requireNonNull(
             trainedModelAllocationNodeService,

+ 2 - 2
x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/inference/deployment/TrainedModelDeploymentTaskTests.java

@@ -18,7 +18,7 @@ import org.elasticsearch.xpack.ml.inference.allocation.TrainedModelAllocationNod
 import java.util.Map;
 import java.util.function.Consumer;
 
-import static org.elasticsearch.xpack.core.ml.MlTasks.TRAINED_MODEL_ALLOCATION_TASK_NAME_PREFIX;
+import static org.elasticsearch.xpack.core.ml.MlTasks.TRAINED_MODEL_ALLOCATION_TASK_ACTION;
 import static org.elasticsearch.xpack.core.ml.MlTasks.TRAINED_MODEL_ALLOCATION_TASK_TYPE;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -32,7 +32,7 @@ public class TrainedModelDeploymentTaskTests extends ESTestCase {
         TrainedModelDeploymentTask task = new TrainedModelDeploymentTask(
             0,
             TRAINED_MODEL_ALLOCATION_TASK_TYPE,
-            TRAINED_MODEL_ALLOCATION_TASK_NAME_PREFIX + modelId,
+            TRAINED_MODEL_ALLOCATION_TASK_ACTION,
             TaskId.EMPTY_TASK_ID,
             Map.of(),
             new StartTrainedModelDeploymentAction.TaskParams(