Add MLTaskrunner and prediction/training transportaction (opensearch-project#37)

zhanghg08 · jackiehanyang · commit 851d49ecbedd · 2021-11-15T19:16:40.000-08:00
* Add MLTaskRunner

* Add MLPredictionTaskRemoteExecutionTransportAction

* Add MLPredictionTaskRemoteExecutionAction

* Add missing HEADER

* Add Training and Prediction Transportaction

* Register transportaction in plugin
diff --git a/plugin/src/main/java/org/opensearch/ml/action/prediction/MLPredictionTaskExecutionAction.java b/plugin/src/main/java/org/opensearch/ml/action/prediction/MLPredictionTaskExecutionAction.java
@@ -0,0 +1,25 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ *
+ */
+
+package org.opensearch.ml.action.prediction;
+
+import org.opensearch.ml.common.transport.prediction.MLPredictionTaskResponse;
+import org.opensearch.action.ActionType;
+
+public class MLPredictionTaskExecutionAction extends ActionType<MLPredictionTaskResponse> {
+    public static MLPredictionTaskExecutionAction INSTANCE = new MLPredictionTaskExecutionAction();
+    public static final String NAME = "cluster:admin/opensearch-ml/prediction/execution";
+
+    private MLPredictionTaskExecutionAction() {
+        super(NAME, MLPredictionTaskResponse::new);
+    }
+}
diff --git a/plugin/src/main/java/org/opensearch/ml/action/prediction/MLPredictionTaskExecutionTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/prediction/MLPredictionTaskExecutionTransportAction.java
@@ -0,0 +1,44 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ *
+ */
+
+package org.opensearch.ml.action.prediction;
+
+import org.opensearch.ml.task.MLTaskRunner;
+import org.opensearch.ml.common.transport.prediction.MLPredictionTaskRequest;
+import org.opensearch.ml.common.transport.prediction.MLPredictionTaskResponse;
+import org.opensearch.action.ActionListener;
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.HandledTransportAction;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.tasks.Task;
+import org.opensearch.transport.TransportService;
+
+public class MLPredictionTaskExecutionTransportAction extends HandledTransportAction<MLPredictionTaskRequest, MLPredictionTaskResponse> {
+    private final MLTaskRunner mlTaskRunner;
+    private final TransportService transportService;
+
+    @Inject
+    public MLPredictionTaskExecutionTransportAction(
+            ActionFilters actionFilters,
+            TransportService transportService,
+            MLTaskRunner mlTaskRunner
+    ) {
+        super(MLPredictionTaskExecutionAction.NAME, transportService, actionFilters, MLPredictionTaskRequest::new);
+        this.mlTaskRunner = mlTaskRunner;
+        this.transportService = transportService;
+    }
+
+    @Override
+    protected void doExecute(Task task, MLPredictionTaskRequest request, ActionListener<MLPredictionTaskResponse> listener) {
+        mlTaskRunner.startPredictionTask(request, listener);
+    }
+}
diff --git a/plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java b/plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java
@@ -0,0 +1,49 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ *
+ */
+
+package org.opensearch.ml.action.prediction;
+
+import lombok.AccessLevel;
+import lombok.experimental.FieldDefaults;
+import lombok.extern.log4j.Log4j2;
+import org.opensearch.action.ActionListener;
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.HandledTransportAction;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.ml.common.transport.prediction.MLPredictionTaskAction;
+import org.opensearch.ml.common.transport.prediction.MLPredictionTaskRequest;
+import org.opensearch.ml.common.transport.prediction.MLPredictionTaskResponse;
+import org.opensearch.ml.task.MLTaskRunner;
+import org.opensearch.tasks.Task;
+import org.opensearch.transport.TransportService;
+
+@Log4j2
+@FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
+public class TransportPredictionTaskAction extends HandledTransportAction<ActionRequest, MLPredictionTaskResponse> {
+    MLTaskRunner mlTaskRunner;
+    TransportService transportService;
+
+    @Inject
+    public TransportPredictionTaskAction(TransportService transportService, ActionFilters actionFilters, MLTaskRunner mlTaskRunner) {
+        super(MLPredictionTaskAction.NAME, transportService, actionFilters, MLPredictionTaskRequest::new);
+        this.mlTaskRunner = mlTaskRunner;
+        this.transportService = transportService;
+    }
+
+    @Override
+    protected void doExecute(Task task, ActionRequest request,
+                             ActionListener<MLPredictionTaskResponse> listener) {
+        MLPredictionTaskRequest mlPredictionTaskRequest = MLPredictionTaskRequest.fromActionRequest(request);
+        mlTaskRunner.runPrediction(mlPredictionTaskRequest, transportService, listener);
+    }
+}
diff --git a/plugin/src/main/java/org/opensearch/ml/action/training/MLTrainingTaskExecutionAction.java b/plugin/src/main/java/org/opensearch/ml/action/training/MLTrainingTaskExecutionAction.java
@@ -0,0 +1,26 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ *
+ */
+
+package org.opensearch.ml.action.training;
+
+import org.opensearch.ml.common.transport.training.MLTrainingTaskResponse;
+import org.opensearch.action.ActionType;
+
+public class MLTrainingTaskExecutionAction extends ActionType<MLTrainingTaskResponse> {
+
+    public static final MLTrainingTaskExecutionAction INSTANCE = new MLTrainingTaskExecutionAction();
+    public static final String NAME = "cluster:admin/opensearch-ml/training/execution";
+
+    public MLTrainingTaskExecutionAction() {
+        super(NAME, MLTrainingTaskResponse::new);
+    }
+}
diff --git a/plugin/src/main/java/org/opensearch/ml/action/training/MLTrainingTaskExecutionTransportAction.java b/plugin/src/main/java/org/opensearch/ml/action/training/MLTrainingTaskExecutionTransportAction.java
@@ -0,0 +1,44 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ *
+ */
+
+package org.opensearch.ml.action.training;
+
+import org.opensearch.ml.common.transport.training.MLTrainingTaskRequest;
+import org.opensearch.ml.common.transport.training.MLTrainingTaskResponse;
+import lombok.AccessLevel;
+import lombok.experimental.FieldDefaults;
+import lombok.extern.log4j.Log4j2;
+import org.opensearch.action.ActionListener;
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.HandledTransportAction;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.ml.task.MLTaskRunner;
+import org.opensearch.tasks.Task;
+import org.opensearch.transport.TransportService;
+
+@Log4j2
+@FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
+public class MLTrainingTaskExecutionTransportAction extends HandledTransportAction<MLTrainingTaskRequest, MLTrainingTaskResponse> {
+    MLTaskRunner mlTaskRunner;
+
+    @Inject
+    public MLTrainingTaskExecutionTransportAction(TransportService transportService, ActionFilters actionFilters,
+                                                  MLTaskRunner mlTaskRunner) {
+        super(MLTrainingTaskExecutionAction.NAME, transportService, actionFilters, MLTrainingTaskRequest::new);
+        this.mlTaskRunner = mlTaskRunner;
+    }
+
+    @Override
+    protected void doExecute(Task task, MLTrainingTaskRequest request, ActionListener<MLTrainingTaskResponse> listener) {
+        mlTaskRunner.startTrainingTask(request, listener);
+    }
+}
diff --git a/plugin/src/main/java/org/opensearch/ml/action/training/TransportTrainingTaskAction.java b/plugin/src/main/java/org/opensearch/ml/action/training/TransportTrainingTaskAction.java
@@ -0,0 +1,47 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ *
+ */
+
+package org.opensearch.ml.action.training;
+
+import lombok.extern.log4j.Log4j2;
+import org.opensearch.action.ActionListener;
+import org.opensearch.action.ActionRequest;
+import org.opensearch.action.support.ActionFilters;
+import org.opensearch.action.support.HandledTransportAction;
+import org.opensearch.common.inject.Inject;
+import org.opensearch.ml.common.transport.training.MLTrainingTaskAction;
+import org.opensearch.ml.common.transport.training.MLTrainingTaskRequest;
+import org.opensearch.ml.common.transport.training.MLTrainingTaskResponse;
+import org.opensearch.ml.task.MLTaskRunner;
+import org.opensearch.tasks.Task;
+import org.opensearch.transport.TransportService;
+
+@Log4j2
+public class TransportTrainingTaskAction extends HandledTransportAction<ActionRequest, MLTrainingTaskResponse> {
+
+    MLTaskRunner mlTaskRunner;
+    TransportService transportService;
+
+    @Inject
+    public TransportTrainingTaskAction(TransportService transportService, ActionFilters actionFilters,
+                                       MLTaskRunner mlTaskRunner) {
+        super(MLTrainingTaskAction.NAME, transportService, actionFilters, MLTrainingTaskRequest::new);
+        this.mlTaskRunner = mlTaskRunner;
+        this.transportService = transportService;
+    }
+
+    @Override
+    protected void doExecute(Task task, ActionRequest request, ActionListener<MLTrainingTaskResponse> listener) {
+        MLTrainingTaskRequest trainingRequest = MLTrainingTaskRequest.fromActionRequest(request);
+        mlTaskRunner.runTraining(trainingRequest, transportService, listener);
+    }
+}
diff --git a/plugin/src/main/java/org/opensearch/ml/indices/MLIndicesHandler.java b/plugin/src/main/java/org/opensearch/ml/indices/MLIndicesHandler.java
@@ -16,8 +16,6 @@
 import lombok.RequiredArgsConstructor;
 import lombok.experimental.FieldDefaults;
 import lombok.extern.log4j.Log4j2;
-import lombok.val;
-import org.opensearch.action.admin.indices.create.CreateIndexRequest;
 import org.opensearch.client.Client;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.xcontent.XContentType;
@@ -31,7 +29,9 @@ public class MLIndicesHandler {
             "    \"properties\": {\n" +
             "      \"taskId\": { \"type\": \"keyword\" },\n" +
             "      \"algorithm\": {\"type\": \"keyword\"},\n" +
-            "      \"model\" : { \"type\": \"binary\"}\n" +
+            "      \"modelName\" : { \"type\": \"keyword\"},\n" +
+            "      \"modelVersion\" : { \"type\": \"keyword\"},\n" +
+            "      \"modelContent\" : { \"type\": \"binary\"}\n" +
             "    }\n" +
             "}";
 
diff --git a/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java b/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java
@@ -28,8 +28,16 @@
 import org.opensearch.common.xcontent.NamedXContentRegistry;
 import org.opensearch.env.Environment;
 import org.opensearch.env.NodeEnvironment;
+import org.opensearch.ml.action.prediction.MLPredictionTaskExecutionAction;
+import org.opensearch.ml.action.prediction.MLPredictionTaskExecutionTransportAction;
+import org.opensearch.ml.action.prediction.TransportPredictionTaskAction;
 import org.opensearch.ml.action.stats.MLStatsNodesAction;
 import org.opensearch.ml.action.stats.MLStatsNodesTransportAction;
+import org.opensearch.ml.action.training.MLTrainingTaskExecutionAction;
+import org.opensearch.ml.action.training.MLTrainingTaskExecutionTransportAction;
+import org.opensearch.ml.action.training.TransportTrainingTaskAction;
+import org.opensearch.ml.common.transport.prediction.MLPredictionTaskAction;
+import org.opensearch.ml.common.transport.training.MLTrainingTaskAction;
 import org.opensearch.ml.rest.RestStatsMLAction;
 import org.opensearch.ml.stats.MLStat;
 import org.opensearch.ml.stats.MLStats;
@@ -44,15 +52,19 @@
 import org.opensearch.rest.RestController;
 import org.opensearch.rest.RestHandler;
 import org.opensearch.script.ScriptService;
+import org.opensearch.threadpool.ExecutorBuilder;
+import org.opensearch.threadpool.FixedExecutorBuilder;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.watcher.ResourceWatcherService;
 
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Supplier;
 
 public class MachineLearningPlugin extends Plugin implements ActionPlugin {
+    public static final String TASK_THREAD_POOL = "OPENSEARCH_ML_TASK_THREAD_POOL";
     public static final String ML_BASE_URI = "/_opensearch/_ml";
 
     private MLStats mlStats;
@@ -72,7 +84,12 @@ public Setting<Boolean> legacySetting() {
     public List<ActionHandler<? extends ActionRequest, ? extends ActionResponse>> getActions() {
         return ImmutableList.of(
                 new ActionHandler<>(MLStatsNodesAction.INSTANCE,
-                        MLStatsNodesTransportAction.class)
+                        MLStatsNodesTransportAction.class),
+                new ActionHandler<>(MLPredictionTaskAction.INSTANCE, TransportPredictionTaskAction.class),
+                new ActionHandler<>(MLTrainingTaskAction.INSTANCE, TransportTrainingTaskAction.class),
+                new ActionHandler<>(MLPredictionTaskExecutionAction.INSTANCE,
+                        MLPredictionTaskExecutionTransportAction.class),
+                new ActionHandler<>(MLTrainingTaskExecutionAction.INSTANCE, MLTrainingTaskExecutionTransportAction.class)
         );
     }
 
@@ -109,4 +126,18 @@ public List<RestHandler> getRestHandlers(
                         restStatsMLAction
                 );
     }
+
+    @Override
+    public List<ExecutorBuilder<?>> getExecutorBuilders(Settings settings) {
+        FixedExecutorBuilder ml = new FixedExecutorBuilder(
+                settings,
+                TASK_THREAD_POOL,
+                4,
+                4,
+                "ml.task_thread_pool",
+                false
+        );
+
+        return Collections.singletonList(ml);
+    }
 }
diff --git a/plugin/src/main/java/org/opensearch/ml/task/MLTaskRunner.java b/plugin/src/main/java/org/opensearch/ml/task/MLTaskRunner.java
