Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Enhance batch job task management by adding default action types #3080

Merged
merged 4 commits into from
Oct 9, 2024
Merged

Enhance batch job task management by adding default action types #3080

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5556080
enhance batch job task management by adding default action types
rbhavna Oct 8, 2024
17cbad0
add unit tests
rbhavna Oct 9, 2024
0bdc387
address comments
rbhavna Oct 9, 2024
14a1188
add default switch case
rbhavna Oct 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions common/src/main/java/org/opensearch/ml/common/connector/AbstractConnector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,11 @@ public Optional<ConnectorAction> findAction(String action) {
return Optional.empty();
}

@Override
public void addAction(ConnectorAction action) {
actions.add(action);
}

@Override
public void removeCredential() {
this.credential = null;
Expand Down
2 changes: 2 additions & 0 deletions common/src/main/java/org/opensearch/ml/common/connector/Connector.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,8 @@ public interface Connector extends ToXContentObject, Writeable {

List<ConnectorAction> getActions();

void addAction(ConnectorAction action);

ConnectorClientConfig getConnectorClientConfig();

String getActionEndpoint(String action, Map<String, String> parameters);
Expand Down
26 changes: 26 additions & 0 deletions common/src/main/java/org/opensearch/ml/common/output/MLPredictionOutput.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
package org.opensearch.ml.common.output;

import java.io.IOException;
import java.util.Map;

import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
Expand All @@ -30,8 +31,12 @@ public class MLPredictionOutput extends MLOutput {
public static final String STATUS_FIELD = "status";
public static final String PREDICTION_RESULT_FIELD = "prediction_result";

// This field will be created for offline batch prediction tasks containing details of the batch job as outputted by the remote server.
public static final String REMOTE_JOB_FIELD = "remote_job";
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's add a comment about this field.


String taskId;
String status;
Map<String, Object> remoteJob;

@ToString.Exclude
DataFrame predictionResult;
Expand All @@ -44,6 +49,14 @@ public MLPredictionOutput(String taskId, String status, DataFrame predictionResu
this.predictionResult = predictionResult;
}

@Builder
public MLPredictionOutput(String taskId, String status, Map<String, Object> remoteJob) {
super(OUTPUT_TYPE);
this.taskId = taskId;
this.status = status;
this.remoteJob = remoteJob;
}

public MLPredictionOutput(StreamInput in) throws IOException {
super(OUTPUT_TYPE);
this.taskId = in.readOptionalString();
Expand All @@ -56,6 +69,9 @@ public MLPredictionOutput(StreamInput in) throws IOException {
break;
}
}
if (in.readBoolean()) {
this.remoteJob = in.readMap(s -> s.readString(), s -> s.readGenericValue());
}
}

@Override
Expand All @@ -69,6 +85,12 @@ public void writeTo(StreamOutput out) throws IOException {
} else {
out.writeBoolean(false);
}
if (remoteJob != null) {
out.writeBoolean(true);
out.writeMap(remoteJob, StreamOutput::writeString, StreamOutput::writeGenericValue);
} else {
out.writeBoolean(false);
}
}

@Override
Expand All @@ -87,6 +109,10 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
builder.endObject();
}

if (remoteJob != null) {
builder.field(REMOTE_JOB_FIELD, remoteJob);
}

builder.endObject();
return builder;
}
Expand Down
14 changes: 14 additions & 0 deletions common/src/test/java/org/opensearch/ml/common/output/MLPredictionOutputTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,9 @@

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.junit.Before;
import org.junit.Test;
Expand All @@ -30,6 +32,7 @@
public class MLPredictionOutputTest {

MLPredictionOutput output;
MLPredictionOutput outputWithRemoteJob;

@Before
public void setUp() {
Expand All @@ -38,12 +41,17 @@ public void setUp() {
rows.add(new Row(new ColumnValue[] { new IntValue(1) }));
rows.add(new Row(new ColumnValue[] { new IntValue(2) }));
DataFrame dataFrame = new DefaultDataFrame(columnMetas, rows);
Map<String, Object> remoteJob = new HashMap<>();
remoteJob.put("status", "INPROGRESS");
remoteJob.put("job_id", "testJobID");
output = MLPredictionOutput.builder().taskId("test_task_id").status("test_status").predictionResult(dataFrame).build();
outputWithRemoteJob = new MLPredictionOutput("test_task_id", "test_status", remoteJob);
}

@Test
public void toXContent() throws IOException {
XContentBuilder builder = MediaTypeRegistry.contentBuilder(XContentType.JSON);
XContentBuilder builderWithRemoteJob = MediaTypeRegistry.contentBuilder(XContentType.JSON);
output.toXContent(builder, ToXContent.EMPTY_PARAMS);
String jsonStr = builder.toString();
assertEquals(
Expand All @@ -53,6 +61,12 @@ public void toXContent() throws IOException {
+ "\"value\":2}]}]}}",
jsonStr
);
outputWithRemoteJob.toXContent(builderWithRemoteJob, ToXContent.EMPTY_PARAMS);
String jsonStr2 = builderWithRemoteJob.toString();
assertEquals(
"{\"task_id\":\"test_task_id\",\"status\":\"test_status\",\"remote_job\":{\"job_id\":\"testJobID\",\"status\":\"INPROGRESS\"}}",
jsonStr2
);
}

@Test
Expand Down
59 changes: 59 additions & 0 deletions ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/ConnectorUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
package org.opensearch.ml.engine.algorithms.remote;

import static org.apache.commons.text.StringEscapeUtils.escapeJson;
import static org.opensearch.ml.common.connector.ConnectorAction.ActionType.BATCH_PREDICT;
import static org.opensearch.ml.common.connector.ConnectorAction.ActionType.CANCEL_BATCH_PREDICT;
import static org.opensearch.ml.common.connector.HttpConnector.RESPONSE_FILTER_FIELD;
import static org.opensearch.ml.common.connector.MLPreProcessFunction.CONVERT_INPUT_TO_JSON_STRING;
Expand All @@ -19,6 +20,7 @@
import java.net.URI;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
Expand Down Expand Up @@ -61,6 +63,9 @@ public class ConnectorUtils {
private static final Aws4Signer signer;
public static final String SKIP_VALIDATE_MISSING_PARAMETERS = "skip_validating_missing_parameters";

public static final List<String> SUPPORTED_REMOTE_SERVERS_FOR_DEFAULT_ACTION_TYPES = List
.of("sagemaker", "openai", "bedrock", "cohere");

Comment on lines +66 to +68
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ideally if a new platform is used but not listed here, CX should still be able to GetTask and CancelTask by manually adding the actions in the connector. But seems this is not the case in this PR?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes if the connector already has the actions configured, then they can get/cancel task for any platform. Only if no action is provided, they we perform this check

static {
signer = Aws4Signer.create();
}
Expand Down Expand Up @@ -313,4 +318,58 @@ public static SdkHttpFullRequest buildSdkRequest(
}
return builder.build();
}

public static ConnectorAction createConnectorAction(Connector connector, ConnectorAction.ActionType actionType) {
Optional<ConnectorAction> batchPredictAction = connector.findAction(BATCH_PREDICT.name());
String predictEndpoint = batchPredictAction.get().getUrl();
Map<String, String> parameters = connector.getParameters() != null
? new HashMap<>(connector.getParameters())
: Collections.emptyMap();

// Apply parameter substitution only if needed
if (!parameters.isEmpty()) {
StringSubstitutor substitutor = new StringSubstitutor(parameters, "${parameters.", "}");
predictEndpoint = substitutor.replace(predictEndpoint);
}

boolean isCancelAction = actionType == CANCEL_BATCH_PREDICT;

// Initialize the default method and requestBody
String method = "POST";
String requestBody = null;
String url = "";

switch (getRemoteServerFromURL(predictEndpoint)) {
case "sagemaker":
url = isCancelAction
? predictEndpoint.replace("CreateTransformJob", "StopTransformJob")
: predictEndpoint.replace("CreateTransformJob", "DescribeTransformJob");
requestBody = "{ \"TransformJobName\" : \"${parameters.TransformJobName}\"}";
break;
case "openai":
case "cohere":
url = isCancelAction ? predictEndpoint + "/${parameters.id}/cancel" : predictEndpoint + "/${parameters.id}";
method = isCancelAction ? "POST" : "GET";
break;
case "bedrock":
url = isCancelAction
? predictEndpoint + "/${parameters.processedJobArn}/stop"
: predictEndpoint + "/${parameters.processedJobArn}";
method = isCancelAction ? "POST" : "GET";
break;
Comment on lines +343 to +359
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should add the default branch for this switch statement to return null. In the GetTask and CancelTask, if the ConnectorAction is null, throw an exception with meaning logs like "please provide GetTask/CancelTask action in the connector".

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh sorry I added it but I guess it got missed during refactoring. Let me add it

}

return ConnectorAction
.builder()
.actionType(actionType)
.method(method)
.url(url)
.requestBody(requestBody)
.headers(batchPredictAction.get().getHeaders())
.build();
}

public static String getRemoteServerFromURL(String url) {
return SUPPORTED_REMOTE_SERVERS_FOR_DEFAULT_ACTION_TYPES.stream().filter(url::contains).findFirst().orElse("");
}
}
Loading
Loading