Adds Register and Deploy Model Step for remote model

build.gradle

@@ -112,6 +112,9 @@ publishing {
 allprojects {
     group = opensearch_group
     version = "${opensearch_build}"
+}
+
+java {
     targetCompatibility = JavaVersion.VERSION_11
     sourceCompatibility = JavaVersion.VERSION_11
 }

gradle.properties

@@ -0,0 +1,38 @@
+#
+# 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.
+#
+
+
+# Enable build caching
+org.gradle.caching=true
+org.gradle.warning.mode=none
+org.gradle.parallel=true
+# Workaround for https://github.com/diffplug/spotless/issues/834
+org.gradle.jvmargs=-Xmx3g -XX:+HeapDumpOnOutOfMemoryError -Xss2m \
+    --add-exports jdk.compiler/com.sun.tools.javac.api=ALL-UNNAMED \
+    --add-exports jdk.compiler/com.sun.tools.javac.file=ALL-UNNAMED \
+    --add-exports jdk.compiler/com.sun.tools.javac.parser=ALL-UNNAMED \
+    --add-exports jdk.compiler/com.sun.tools.javac.tree=ALL-UNNAMED \
+    --add-exports jdk.compiler/com.sun.tools.javac.util=ALL-UNNAMED
+options.forkOptions.memoryMaximumSize=2g
+
+# Disable duplicate project id detection
+# See https://docs.gradle.org/current/userguide/upgrading_version_6.html#duplicate_project_names_may_cause_publication_to_fail
+systemProp.org.gradle.dependency.duplicate.project.detection=false
+
+# Enforce the build to fail on deprecated gradle api usage
+systemProp.org.gradle.warning.mode=fail
+
+# forcing to use TLS1.2 to avoid failure in vault
+# see https://github.com/hashicorp/vault/issues/8750#issuecomment-631236121
+systemProp.jdk.tls.client.protocols=TLSv1.2
+
+# jvm args for faster test execution by default
+systemProp.tests.jvm.argline=-XX:TieredStopAtLevel=1 -XX:ReservedCodeCacheSize=64m

src/main/java/demo/DemoWorkflowStep.java

@@ -37,7 +37,7 @@ public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
         CompletableFuture.runAsync(() -> {
             try {
                 Thread.sleep(this.delay);
-                future.complete(null);
+                future.complete(WorkflowData.EMPTY);
             } catch (InterruptedException e) {
                 future.completeExceptionally(e);
             }

src/main/java/demo/TemplateParseDemo.java

@@ -55,7 +55,6 @@ public static void main(String[] args) throws IOException {
         }
         ClusterService clusterService = new ClusterService(null, null, null);
         Client client = new NodeClient(null, null);
-
         WorkflowStepFactory factory = new WorkflowStepFactory(clusterService, client);
         ThreadPool threadPool = new ThreadPool(Settings.EMPTY);
         WorkflowProcessSorter workflowProcessSorter = new WorkflowProcessSorter(factory, threadPool);

src/main/java/org/opensearch/flowframework/client/MLClient.java

@@ -8,7 +8,7 @@
  */
 package org.opensearch.flowframework.client;
 
-import org.opensearch.client.node.NodeClient;
+import org.opensearch.client.Client;
 import org.opensearch.ml.client.MachineLearningNodeClient;
 
 /**
@@ -22,12 +22,12 @@ private MLClient() {}
     /**
      * Creates machine learning client.
      *
-     * @param nodeClient node client of OpenSearch.
+     * @param client client of OpenSearch.
      * @return machine learning client from ml-commons.
      */
-    public static MachineLearningNodeClient createMLClient(NodeClient nodeClient) {
+    public static MachineLearningNodeClient createMLClient(Client client) {
         if (INSTANCE == null) {
-            INSTANCE = new MachineLearningNodeClient(nodeClient);
+            INSTANCE = new MachineLearningNodeClient(client);
         }
         return INSTANCE;
     }

src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java

@@ -0,0 +1,83 @@
+/*
+ * Copyright OpenSearch Contributors
+ * 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.
+ */
+package org.opensearch.flowframework.workflow;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.client.Client;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.flowframework.client.MLClient;
+import org.opensearch.ml.client.MachineLearningNodeClient;
+import org.opensearch.ml.common.transport.deploy.MLDeployModelResponse;
+
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * Step to deploy a model
+ */
+public class DeployModelStep implements WorkflowStep {
+    private static final Logger logger = LogManager.getLogger(DeployModelStep.class);
+
+    private Client client;
+    private static final String MODEL_ID = "model_id";
+    static final String NAME = "deploy_model";
+
+    /**
+     * Instantiate this class
+     * @param client client to instantiate MLClient
+     */
+    public DeployModelStep(Client client) {
+        this.client = client;
+    }
+
+    @Override
+    public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
+
+        CompletableFuture<WorkflowData> deployModelFuture = new CompletableFuture<>();
+
+        MachineLearningNodeClient machineLearningNodeClient = MLClient.createMLClient(client);
+
+        ActionListener<MLDeployModelResponse> actionListener = new ActionListener<>() {
+            @Override
+            public void onResponse(MLDeployModelResponse mlDeployModelResponse) {
+                logger.info("Model deployment state {}", mlDeployModelResponse.getStatus());
+                deployModelFuture.complete(
+                    new WorkflowData(Map.ofEntries(Map.entry("deploy_model_status", mlDeployModelResponse.getStatus())))
+                );
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                logger.error("Model deployment failed");
+                deployModelFuture.completeExceptionally(e);
+            }
+        };
+
+        String modelId = null;
+
+        for (WorkflowData workflowData : data) {
+            Map<String, Object> content = workflowData.getContent();
+            for (Map.Entry<String, Object> entry : content.entrySet()) {
+                if (entry.getKey() == MODEL_ID) {
+                    modelId = (String) content.get(MODEL_ID);
+                }
+
+            }
+        }
+        machineLearningNodeClient.deploy(modelId, actionListener);
+        return deployModelFuture;
+    }
+
+    @Override
+    public String getName() {
+        return NAME;
+    }
+}

src/main/java/org/opensearch/flowframework/workflow/GetTask.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright OpenSearch Contributors
+ * 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.
+ */
+package org.opensearch.flowframework.workflow;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.common.SuppressForbidden;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.ml.client.MachineLearningNodeClient;
+import org.opensearch.ml.common.MLTask;
+import org.opensearch.ml.common.MLTaskState;
+import org.opensearch.ml.common.transport.task.MLTaskGetResponse;
+
+/**
+ * Step to get modelID of a registered local model
+ */
+@SuppressForbidden(reason = "This class is for the future work of registering local model")
+public class GetTask {
+
+    private static final Logger logger = LogManager.getLogger(GetTask.class);
+    private MachineLearningNodeClient machineLearningNodeClient;
+    private String taskId;
+
+    /**
+     * Instantiate this class
+     * @param machineLearningNodeClient client to instantiate ml-commons APIs
+     * @param taskId taskID of the model
+     */
+    public GetTask(MachineLearningNodeClient machineLearningNodeClient, String taskId) {
+        this.machineLearningNodeClient = machineLearningNodeClient;
+        this.taskId = taskId;
+    }
+
+    /**
+     * Invokes get task API of ml-commons
+     */
+    public void getTask() {
+
+        ActionListener<MLTask> actionListener = new ActionListener<>() {
+            @Override
+            public void onResponse(MLTask mlTask) {
+                if (mlTask.getState() == MLTaskState.COMPLETED) {
+                    logger.info("Model registration successful");
+                    MLTaskGetResponse response = MLTaskGetResponse.builder().mlTask(mlTask).build();
+                    logger.info("Response from task {}", response);
+                }
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                logger.error("Model registration failed");
+            }
+        };
+
+        machineLearningNodeClient.getTask(taskId, actionListener);
+
+    }
+
+}

src/main/java/org/opensearch/flowframework/workflow/RegisterModelStep.java

@@ -0,0 +1,152 @@
+/*
+ * Copyright OpenSearch Contributors
+ * 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.
+ */
+package org.opensearch.flowframework.workflow;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.client.Client;
+import org.opensearch.core.action.ActionListener;
+import org.opensearch.flowframework.client.MLClient;
+import org.opensearch.ml.client.MachineLearningNodeClient;
+import org.opensearch.ml.common.FunctionName;
+import org.opensearch.ml.common.model.MLModelConfig;
+import org.opensearch.ml.common.model.MLModelFormat;
+import org.opensearch.ml.common.transport.register.MLRegisterModelInput;
+import org.opensearch.ml.common.transport.register.MLRegisterModelResponse;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Stream;
+
+/**
+ * Step to register a remote model
+ */
+public class RegisterModelStep implements WorkflowStep {
+
+    private static final Logger logger = LogManager.getLogger(RegisterModelStep.class);
+
+    private Client client;
+
+    static final String NAME = "register_model";
+
+    private static final String FUNCTION_NAME = "function_name";
+    private static final String MODEL_NAME = "name";
+    private static final String MODEL_VERSION = "model_version";
+    private static final String MODEL_GROUP_ID = "model_group_id";
+    private static final String DESCRIPTION = "description";
+    private static final String CONNECTOR_ID = "connector_id";
+    private static final String MODEL_FORMAT = "model_format";
+    private static final String MODEL_CONFIG = "model_config";
+
+    /**
+     * Instantiate this class
+     * @param client client to instantiate MLClient
+     */
+    public RegisterModelStep(Client client) {
+        this.client = client;
+    }
+
+    @Override
+    public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
+
+        CompletableFuture<WorkflowData> registerModelFuture = new CompletableFuture<>();
+
+        MachineLearningNodeClient machineLearningNodeClient = MLClient.createMLClient(client);
+
+        ActionListener<MLRegisterModelResponse> actionListener = new ActionListener<>() {
+            @Override
+            public void onResponse(MLRegisterModelResponse mlRegisterModelResponse) {
+                logger.info("Model registration successful");
+                registerModelFuture.complete(
+                    new WorkflowData(
+                        Map.ofEntries(
+                            Map.entry("model_id", mlRegisterModelResponse.getModelId()),
+                            Map.entry("register_model_status", mlRegisterModelResponse.getStatus())
+                        )
+                    )
+                );
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                logger.error("Failed to register model");
+                registerModelFuture.completeExceptionally(new IOException("Failed to register model "));
+            }
+        };
+
+        FunctionName functionName = null;
+        String modelName = null;
+        String modelVersion = null;
+        String modelGroupId = null;
+        String connectorId = null;
+        String description = null;
+        MLModelFormat modelFormat = null;
+        MLModelConfig modelConfig = null;
+
+        for (WorkflowData workflowData : data) {
+            Map<String, Object> content = workflowData.getContent();
+
+            for (Entry<String, Object> entry : content.entrySet()) {
+                switch (entry.getKey()) {
+                    case FUNCTION_NAME:
+                        functionName = FunctionName.from(((String) content.get(FUNCTION_NAME)).toUpperCase(Locale.ROOT));
+                        break;
+                    case MODEL_NAME:
+                        modelName = (String) content.get(MODEL_NAME);
+                        break;
+                    case MODEL_VERSION:
+                        modelVersion = (String) content.get(MODEL_VERSION);
+                        break;
+                    case MODEL_GROUP_ID:
+                        modelGroupId = (String) content.get(MODEL_GROUP_ID);
+                        break;
+                    case MODEL_FORMAT:
+                        modelFormat = MLModelFormat.from((String) content.get(MODEL_FORMAT));
+                        break;
+                    case MODEL_CONFIG:
+                        modelConfig = (MLModelConfig) content.get(MODEL_CONFIG);
+                        break;
+                    case DESCRIPTION:
+                        description = (String) content.get(DESCRIPTION);
+                        break;
+                    case CONNECTOR_ID:
+                        connectorId = (String) content.get(CONNECTOR_ID);
+                        break;
+                    default:
+                        break;
+
+                }
+            }
+        }
+
+        if (Stream.of(functionName, modelName, description, connectorId).allMatch(x -> x != null)) {
+
+            // TODO: Add model Config and type cast correctly
+            MLRegisterModelInput mlInput = MLRegisterModelInput.builder()
+                .functionName(functionName)
+                .modelName(modelName)
+                .description(description)
+                .connectorId(connectorId)
+                .build();
+
+            machineLearningNodeClient.register(mlInput, actionListener);
+        }
+
+        return registerModelFuture;
+    }
+
+    @Override
+    public String getName() {
+        return NAME;
+    }
+}