Skip to content

Commit

Permalink
Rebase
Browse files Browse the repository at this point in the history
Signed-off-by: Joshua Palis <[email protected]>
  • Loading branch information
joshpalis committed Nov 30, 2023
2 parents 1417ba4 + 7c06865 commit b51d3ef
Show file tree
Hide file tree
Showing 22 changed files with 345 additions and 80 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down Expand Up @@ -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()) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<>();

Expand All @@ -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) {

Expand Down Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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<>();

Expand All @@ -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()
)
);
}
Expand All @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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<>();

Expand All @@ -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()
)
);
}
Expand All @@ -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();

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

/**
Expand All @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import org.opensearch.ml.common.transport.register.MLRegisterModelInput.MLRegisterModelInputBuilder;
import org.opensearch.ml.common.transport.register.MLRegisterModelResponse;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
Expand Down Expand Up @@ -71,20 +72,30 @@ public RegisterLocalModelStep(Settings settings, ClusterService clusterService,
}

@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> registerLocalModelFuture = new CompletableFuture<>();

// 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());

ActionListener<MLRegisterModelResponse> actionListener = new ActionListener<>() {
@Override
public void onResponse(MLRegisterModelResponse mlRegisterModelResponse) {
logger.info("Local Model registration task creation successful");

String workflowId = data.get(0).getWorkflowId();
String taskId = mlRegisterModelResponse.getTaskId();

// Attempt to retrieve the model ID
retryableGetMlTask(workflowId, registerLocalModelFuture, taskId, 0);
retryableGetMlTask(currentNodeInputs.getWorkflowId(), currentNodeId, registerLocalModelFuture, taskId, 0);
}

@Override
Expand Down Expand Up @@ -206,11 +217,18 @@ public String getName() {
/**
* Retryable get ml task
* @param workflowId the workflow id
* @param nodeId the workflow node id
* @param getMLTaskFuture the workflow step future
* @param taskId the ml task id
* @param retries the current number of request retries
*/
void retryableGetMlTask(String workflowId, CompletableFuture<WorkflowData> registerLocalModelFuture, String taskId, int retries) {
void retryableGetMlTask(
String workflowId,
String nodeId,
CompletableFuture<WorkflowData> registerLocalModelFuture,
String taskId,
int retries
) {
mlClient.getTask(taskId, ActionListener.wrap(response -> {
MLTaskState currentState = response.getState();
if (currentState != MLTaskState.COMPLETED) {
Expand All @@ -231,7 +249,8 @@ void retryableGetMlTask(String workflowId, CompletableFuture<WorkflowData> regis
Map.entry(MODEL_ID, response.getModelId()),
Map.entry(REGISTER_MODEL_STATUS, response.getState().name())
),
workflowId
workflowId,
nodeId
)
);
}
Expand All @@ -244,7 +263,7 @@ void retryableGetMlTask(String workflowId, CompletableFuture<WorkflowData> regis
FutureUtils.cancel(registerLocalModelFuture);
}
final int retryAdd = retries + 1;
retryableGetMlTask(workflowId, registerLocalModelFuture, taskId, retryAdd);
retryableGetMlTask(workflowId, nodeId, registerLocalModelFuture, taskId, retryAdd);
} else {
logger.error("Failed to retrieve local model registration task, maximum retries exceeded");
registerLocalModelFuture.completeExceptionally(
Expand Down
Loading

0 comments on commit b51d3ef

Please sign in to comment.