[Backport-main] Improve workflow step execute API (#215) (#221)

Improve workflow step execute API (#215) * Improve workflow step execute API * Fix typo --------- Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Nov 30, 2023 · 7c06865 · 7c06865
1 parent 772fbbd
commit 7c06865
Show file tree

Hide file tree

Showing 24 changed files with 373 additions and 85 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/workflow/CreateConnectorStep.java b/src/main/java/org/opensearch/flowframework/workflow/CreateConnectorStep.java
@@ -74,22 +74,28 @@ public CreateConnectorStep(MachineLearningNodeClient mlClient, FlowFrameworkIndi
 
     // TODO: need to add retry conflicts here
     @Override
-    public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) throws IOException {
+    public CompletableFuture<WorkflowData> execute(
+        String currentNodeId,
+        WorkflowData currentNodeInputs,
+        Map<String, WorkflowData> outputs,
+        Map<String, String> previousNodeInputs
+    ) throws IOException {
         CompletableFuture<WorkflowData> createConnectorFuture = new CompletableFuture<>();
 
         ActionListener<MLCreateConnectorResponse> actionListener = new ActionListener<>() {
 
             @Override
             public void onResponse(MLCreateConnectorResponse mlCreateConnectorResponse) {
+                String workflowId = currentNodeInputs.getWorkflowId();
                 createConnectorFuture.complete(
                     new WorkflowData(
                         Map.ofEntries(Map.entry("connector_id", mlCreateConnectorResponse.getConnectorId())),
-                        data.get(0).getWorkflowId()
+                        workflowId,
+                        currentNodeInputs.getNodeId()
                     )
                 );
                 try {
                     logger.info("Created connector successfully");
-                    String workflowId = data.get(0).getWorkflowId();
                     String workflowStepName = getName();
                     ResourceCreated newResource = new ResourceCreated(workflowStepName, mlCreateConnectorResponse.getConnectorId());
                     XContentBuilder builder = XContentBuilder.builder(XContentType.JSON.xContent());
@@ -136,6 +142,12 @@ public void onFailure(Exception e) {
         Map<String, String> credentials = Collections.emptyMap();
         List<ConnectorAction> actions = Collections.emptyList();
 
+        // TODO: Recreating the list to get this compiling
+        // Need to refactor the below iteration to pull directly from the maps
+        List<WorkflowData> data = new ArrayList<>();
+        data.add(currentNodeInputs);
+        data.addAll(outputs.values());
+
         try {
             for (WorkflowData workflowData : data) {
                 for (Entry<String, Object> entry : workflowData.getContent().entrySet()) {

diff --git a/src/main/java/org/opensearch/flowframework/workflow/CreateIndexStep.java b/src/main/java/org/opensearch/flowframework/workflow/CreateIndexStep.java
@@ -19,6 +19,7 @@
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.flowframework.indices.FlowFrameworkIndicesHandler;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -54,14 +55,25 @@ public CreateIndexStep(ClusterService clusterService, Client client) {
     }
 
     @Override
-    public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
+    public CompletableFuture<WorkflowData> execute(
+        String currentNodeId,
+        WorkflowData currentNodeInputs,
+        Map<String, WorkflowData> outputs,
+        Map<String, String> previousNodeInputs
+    ) {
         CompletableFuture<WorkflowData> future = new CompletableFuture<>();
         ActionListener<CreateIndexResponse> actionListener = new ActionListener<>() {
 
             @Override
             public void onResponse(CreateIndexResponse createIndexResponse) {
                 logger.info("created index: {}", createIndexResponse.index());
-                future.complete(new WorkflowData(Map.of(INDEX_NAME, createIndexResponse.index()), data.get(0).getWorkflowId()));
+                future.complete(
+                    new WorkflowData(
+                        Map.of(INDEX_NAME, createIndexResponse.index()),
+                        currentNodeInputs.getWorkflowId(),
+                        currentNodeInputs.getNodeId()
+                    )
+                );
             }
 
             @Override
@@ -75,6 +87,12 @@ public void onFailure(Exception e) {
         String type = null;
         Settings settings = null;
 
+        // TODO: Recreating the list to get this compiling
+        // Need to refactor the below iteration to pull directly from the maps
+        List<WorkflowData> data = new ArrayList<>();
+        data.add(currentNodeInputs);
+        data.addAll(outputs.values());
+
         for (WorkflowData workflowData : data) {
             Map<String, Object> content = workflowData.getContent();
             index = (String) content.get(INDEX_NAME);

diff --git a/src/main/java/org/opensearch/flowframework/workflow/CreateIngestPipelineStep.java b/src/main/java/org/opensearch/flowframework/workflow/CreateIngestPipelineStep.java
@@ -20,6 +20,7 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -59,7 +60,12 @@ public CreateIngestPipelineStep(Client client) {
     }
 
     @Override
-    public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
+    public CompletableFuture<WorkflowData> execute(
+        String currentNodeId,
+        WorkflowData currentNodeInputs,
+        Map<String, WorkflowData> outputs,
+        Map<String, String> previousNodeInputs
+    ) {
 
         CompletableFuture<WorkflowData> createIngestPipelineFuture = new CompletableFuture<>();
 
@@ -71,6 +77,12 @@ public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
         String outputFieldName = null;
         BytesReference configuration = null;
 
+        // TODO: Recreating the list to get this compiling
+        // Need to refactor the below iteration to pull directly from the maps
+        List<WorkflowData> data = new ArrayList<>();
+        data.add(currentNodeInputs);
+        data.addAll(outputs.values());
+
         // Extract required content from workflow data and generate the ingest pipeline configuration
         for (WorkflowData workflowData : data) {
 
@@ -126,7 +138,11 @@ public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
 
                 // PutPipelineRequest returns only an AcknowledgeResponse, returning pipelineId instead
                 createIngestPipelineFuture.complete(
-                    new WorkflowData(Map.of(PIPELINE_ID, putPipelineRequest.getId()), data.get(0).getWorkflowId())
+                    new WorkflowData(
+                        Map.of(PIPELINE_ID, putPipelineRequest.getId()),
+                        currentNodeInputs.getWorkflowId(),
+                        currentNodeInputs.getNodeId()
+                    )
                 );
 
                 // TODO : Use node client to index response data to global context (pending global context index implementation)

diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeployModelStep.java
@@ -17,6 +17,7 @@
 import org.opensearch.ml.client.MachineLearningNodeClient;
 import org.opensearch.ml.common.transport.deploy.MLDeployModelResponse;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
@@ -41,7 +42,12 @@ public DeployModelStep(MachineLearningNodeClient mlClient) {
     }
 
     @Override
-    public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
+    public CompletableFuture<WorkflowData> execute(
+        String currentNodeId,
+        WorkflowData currentNodeInputs,
+        Map<String, WorkflowData> outputs,
+        Map<String, String> previousNodeInputs
+    ) {
 
         CompletableFuture<WorkflowData> deployModelFuture = new CompletableFuture<>();
 
@@ -52,7 +58,8 @@ public void onResponse(MLDeployModelResponse mlDeployModelResponse) {
                 deployModelFuture.complete(
                     new WorkflowData(
                         Map.ofEntries(Map.entry("deploy_model_status", mlDeployModelResponse.getStatus())),
-                        data.get(0).getWorkflowId()
+                        currentNodeInputs.getWorkflowId(),
+                        currentNodeInputs.getNodeId()
                     )
                 );
             }
@@ -66,6 +73,12 @@ public void onFailure(Exception e) {
 
         String modelId = null;
 
+        // TODO: Recreating the list to get this compiling
+        // Need to refactor the below iteration to pull directly from the maps
+        List<WorkflowData> data = new ArrayList<>();
+        data.add(currentNodeInputs);
+        data.addAll(outputs.values());
+
         for (WorkflowData workflowData : data) {
             if (workflowData.getContent().containsKey(MODEL_ID)) {
                 modelId = (String) workflowData.getContent().get(MODEL_ID);

diff --git a/src/main/java/org/opensearch/flowframework/workflow/GetMLTaskStep.java b/src/main/java/org/opensearch/flowframework/workflow/GetMLTaskStep.java
@@ -20,6 +20,7 @@
 import org.opensearch.ml.client.MachineLearningNodeClient;
 import org.opensearch.ml.common.MLTaskState;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -50,12 +51,23 @@ public GetMLTaskStep(Settings settings, ClusterService clusterService, MachineLe
     }
 
     @Override
-    public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
+    public CompletableFuture<WorkflowData> execute(
+        String currentNodeId,
+        WorkflowData currentNodeInputs,
+        Map<String, WorkflowData> outputs,
+        Map<String, String> previousNodeInputs
+    ) {
 
         CompletableFuture<WorkflowData> getMLTaskFuture = new CompletableFuture<>();
 
         String taskId = null;
 
+        // TODO: Recreating the list to get this compiling
+        // Need to refactor the below iteration to pull directly from the maps
+        List<WorkflowData> data = new ArrayList<>();
+        data.add(currentNodeInputs);
+        data.addAll(outputs.values());
+
         for (WorkflowData workflowData : data) {
             Map<String, Object> content = workflowData.getContent();
             for (Entry<String, Object> entry : content.entrySet()) {
@@ -73,7 +85,7 @@ public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) {
             logger.error("Failed to retrieve ML Task");
             getMLTaskFuture.completeExceptionally(new FlowFrameworkException("Required fields are not provided", RestStatus.BAD_REQUEST));
         } else {
-            retryableGetMlTask(data.get(0).getWorkflowId(), getMLTaskFuture, taskId, 0);
+            retryableGetMlTask(currentNodeInputs.getWorkflowId(), currentNodeInputs.getNodeId(), getMLTaskFuture, taskId, 0);
         }
 
         return getMLTaskFuture;
@@ -87,11 +99,18 @@ public String getName() {
     /**
      * Retryable GetMLTask
      * @param workflowId the workflow id
+     * @param nodeId the node id
      * @param getMLTaskFuture the workflow step future
      * @param taskId the ml task id
      * @param retries the current number of request retries
      */
-    protected void retryableGetMlTask(String workflowId, CompletableFuture<WorkflowData> getMLTaskFuture, String taskId, int retries) {
+    protected void retryableGetMlTask(
+        String workflowId,
+        String nodeId,
+        CompletableFuture<WorkflowData> getMLTaskFuture,
+        String taskId,
+        int retries
+    ) {
         mlClient.getTask(taskId, ActionListener.wrap(response -> {
             if (response.getState() != MLTaskState.COMPLETED) {
                 throw new IllegalStateException("MLTask is not yet completed");
@@ -103,7 +122,8 @@ protected void retryableGetMlTask(String workflowId, CompletableFuture<WorkflowD
                             Map.entry(MODEL_ID, response.getModelId()),
                             Map.entry(REGISTER_MODEL_STATUS, response.getState().name())
                         ),
-                        workflowId
+                        workflowId,
+                        nodeId
                     )
                 );
             }
@@ -116,7 +136,7 @@ protected void retryableGetMlTask(String workflowId, CompletableFuture<WorkflowD
                     FutureUtils.cancel(getMLTaskFuture);
                 }
                 final int retryAdd = retries + 1;
-                retryableGetMlTask(workflowId, getMLTaskFuture, taskId, retryAdd);
+                retryableGetMlTask(workflowId, nodeId, getMLTaskFuture, taskId, retryAdd);
             } else {
                 logger.error("Failed to retrieve ML Task, maximum retries exceeded");
                 getMLTaskFuture.completeExceptionally(

diff --git a/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java
@@ -53,7 +53,12 @@ public ModelGroupStep(MachineLearningNodeClient mlClient) {
     }
 
     @Override
-    public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) throws IOException {
+    public CompletableFuture<WorkflowData> execute(
+        String currentNodeId,
+        WorkflowData currentNodeInputs,
+        Map<String, WorkflowData> outputs,
+        Map<String, String> previousNodeInputs
+    ) throws IOException {
 
         CompletableFuture<WorkflowData> registerModelGroupFuture = new CompletableFuture<>();
 
@@ -67,7 +72,8 @@ public void onResponse(MLRegisterModelGroupResponse mlRegisterModelGroupResponse
                             Map.entry("model_group_id", mlRegisterModelGroupResponse.getModelGroupId()),
                             Map.entry("model_group_status", mlRegisterModelGroupResponse.getStatus())
                         ),
-                        data.get(0).getWorkflowId()
+                        currentNodeInputs.getWorkflowId(),
+                        currentNodeInputs.getNodeId()
                     )
                 );
             }
@@ -85,6 +91,12 @@ public void onFailure(Exception e) {
         AccessMode modelAccessMode = null;
         Boolean isAddAllBackendRoles = null;
 
+        // TODO: Recreating the list to get this compiling
+        // Need to refactor the below iteration to pull directly from the maps
+        List<WorkflowData> data = new ArrayList<>();
+        data.add(currentNodeInputs);
+        data.addAll(outputs.values());
+
         for (WorkflowData workflowData : data) {
             Map<String, Object> content = workflowData.getContent();
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/NoOpStep.java b/src/main/java/org/opensearch/flowframework/workflow/NoOpStep.java
@@ -9,7 +9,7 @@
 package org.opensearch.flowframework.workflow;
 
 import java.io.IOException;
-import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -21,7 +21,12 @@ public class NoOpStep implements WorkflowStep {
     public static final String NAME = "noop";
 
     @Override
-    public CompletableFuture<WorkflowData> execute(List<WorkflowData> data) throws IOException {
+    public CompletableFuture<WorkflowData> execute(
+        String currentNodeId,
+        WorkflowData currentNodeInputs,
+        Map<String, WorkflowData> outputs,
+        Map<String, String> previousNodeInputs
+    ) throws IOException {
         return CompletableFuture.completedFuture(WorkflowData.EMPTY);
     }
 

diff --git a/src/main/java/org/opensearch/flowframework/workflow/ProcessNode.java b/src/main/java/org/opensearch/flowframework/workflow/ProcessNode.java
@@ -14,7 +14,7 @@
 import org.opensearch.threadpool.Scheduler.ScheduledCancellable;
 import org.opensearch.threadpool.ThreadPool;
 
-import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
@@ -152,10 +152,10 @@ public CompletableFuture<WorkflowData> execute() {
 
                 logger.info("Starting {}.", this.id);
                 // get the input data from predecessor(s)
-                List<WorkflowData> input = new ArrayList<WorkflowData>();
-                input.add(this.input);
+                Map<String, WorkflowData> inputMap = new HashMap<>();
                 for (CompletableFuture<WorkflowData> cf : predFutures) {
-                    input.add(cf.get());
+                    WorkflowData wd = cf.get();
+                    inputMap.put(wd.getNodeId(), wd);
                 }
 
                 ScheduledCancellable delayExec = null;
@@ -167,7 +167,12 @@ public CompletableFuture<WorkflowData> execute() {
                     }, this.nodeTimeout, ThreadPool.Names.SAME);
                 }
                 // record start time for this step.
-                CompletableFuture<WorkflowData> stepFuture = this.workflowStep.execute(input);
+                CompletableFuture<WorkflowData> stepFuture = this.workflowStep.execute(
+                    this.id,
+                    this.input,
+                    inputMap,
+                    this.previousNodeInputs
+                );
                 // If completed exceptionally, this is a no-op
                 future.complete(stepFuture.get());
                 // record end time passing workflow steps