support get batch transform job status in get task API #2825

rbhavna · 2024-08-13T18:28:46Z

Description

support get batch transform job status in get task API

//TODO: Will be adding UTs and ITs in next revision

Related Issues

Resolves #[Issue number to be closed when this PR is merged]

Check List

New functionality includes testing.
New functionality has been documented.
API changes companion pull request created.
Commits are signed per the DCO using --signoff.
Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

Zhangxunmt · 2024-08-29T23:42:17Z

common/src/main/java/org/opensearch/ml/common/connector/ConnectorAction.java

+        CANCEL_BATCH,
+        BATCH_STATUS;



Suggest rename these two action type to CANCEL_BATCH_PREDICT and BATCH_PREDICT_STATUS for sake of consistency and extensibility. It's possible that other batch action like batch_ingest will be added in the future so only using batch may look confused.

Zhangxunmt · 2024-08-29T23:43:51Z

plugin/src/main/java/org/opensearch/ml/action/tasks/GetTaskTransportAction.java

+    }
+
+    private void processRemoteBatchPrediction(MLTask mlTask, String taskId, ActionListener<MLTaskGetResponse> actionListener) {
+


this empty line seems redundant.

Zhangxunmt · 2024-08-29T23:54:07Z

common/src/main/java/org/opensearch/ml/common/MLTask.java

@@ -275,6 +314,9 @@ public static MLTask parse(XContentParser parser) throws IOException {
                case IS_ASYNC_TASK_FIELD:
                    async = parser.booleanValue();
                    break;
+                case TRANSFORM_JOB_FIELD:
+                    transformJob = parser.map();


later in the transport action. It's required all the values in this transformJob map is a String. Why not use parser.mapString() directly?

No the values of transformJob are objects, not strings

Zhangxunmt · 2024-08-29T23:55:38Z

plugin/src/main/java/org/opensearch/ml/action/tasks/GetTaskTransportAction.java

+        Map<String, String> parameters = new HashMap<>();
+        for (Map.Entry<String, ?> entry : transformJob.entrySet()) {
+            if (entry.getValue() instanceof String) {
+                parameters.put(entry.getKey(), (String) entry.getValue());
+            } else {
+                log.debug("Value for key " + entry.getKey() + " is not a String");
+            }
+        }


This component can be moved to the StringUtils class for better sharing and reusing.

sure will do it

Zhangxunmt · 2024-08-30T00:06:54Z

plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java

+                                            .getDataAsMap();
+                                        if (dataAsMap != null
+                                            && (dataAsMap.containsKey("TransformJobArn") || dataAsMap.containsKey("id"))) {
+                                            transformJob.putAll(dataAsMap);


I think you can just convert all the Object to String for easier data handling in the Get Task action, without any loss.

Okay. Let me try to make that change and verify

Zhangxunmt · 2024-08-30T00:09:33Z

plugin/src/main/java/org/opensearch/ml/action/tasks/GetTaskTransportAction.java

+                    ActionListener<Connector> listener = ActionListener.wrap(connector -> {
+                        connector.decrypt(BATCH_STATUS.name(), (credential) -> encryptor.decrypt(credential));
+                        RemoteConnectorExecutor connectorExecutor = MLEngineClassLoader


Is connector access controller included? If it's just one line change, I think we should still add this access check.

You mean to check the user's access to the connector

We should then also add model access controller I guess

Zhangxunmt · 2024-08-30T00:13:50Z

This would need the new action type "task status" in the connector actions body. Can you pls add one example in the test doc or this PR for this new action type? Also UTs are still missing for the GetTask Transport action.

ylwu-amzn · 2024-08-30T16:59:52Z

common/src/main/java/org/opensearch/ml/common/MLTask.java

@@ -45,6 +51,8 @@ public class MLTask implements ToXContentObject, Writeable {
    public static final String LAST_UPDATE_TIME_FIELD = "last_update_time";
    public static final String ERROR_FIELD = "error";
    public static final String IS_ASYNC_TASK_FIELD = "is_async";
+    public static final String TRANSFORM_JOB_FIELD = "transform_job";


transform_job seems only limited to Sagemaker, how about changing to remote_job? cc @Zhangxunmt

"remote job" sounds good to me.

sure makes sense

How about remote_batch_job or batch_job

I think not all remote jobs are for batching purpose

ylwu-amzn · 2024-08-30T17:07:54Z

common/src/main/java/org/opensearch/ml/common/MLTask.java

+                out.writeBoolean(true);
+                try {
+                    AccessController.doPrivileged((PrivilegedExceptionAction<Void>) () -> {
+                        out.writeString(gson.toJson(transformJob));


Can we use out.writeMap(..) ?

output.writeMap(transformJob, StreamOutput::writeString, StreamOutput::writeGenericValue)

For Map<String, Object> type? I am trying to change transformJob's type to Map<String, String> but getting some exceptions while testing

ylwu-amzn · 2024-08-30T17:12:19Z

plugin/src/main/java/org/opensearch/ml/action/tasks/GetTaskTransportAction.java

+            }
+        }
+
+        if (parameters.containsKey("TransformJobArn") && parameters.get("TransformJobArn") != null) {


Add some comments?
I think TransformJobArn and TransformJobArn only for Sagemaker.

ylwu-amzn · 2024-08-30T17:27:50Z

plugin/src/main/java/org/opensearch/ml/action/tasks/GetTaskTransportAction.java

+                    });
+                    try (ThreadContext.StoredContext threadContext = client.threadPool().getThreadContext().stashContext()) {
+                        connectorAccessControlHelper
+                            .getConnector(client, model.getConnectorId(), ActionListener.runBefore(listener, threadContext::restore));


Should we also consider internal connector in model? There will be no connector id

Thanks for the catch. Let me add that case too

ylwu-amzn · 2024-08-30T21:09:06Z

plugin/src/main/java/org/opensearch/ml/action/tasks/GetTaskTransportAction.java

+                        updatedTask.put(TRANSFORM_JOB_FIELD, transformJob);
+
+                        if ((transformJob.containsKey("status") && transformJob.get("status").equals("completed"))
+                            || (transformJob.containsKey("TransformJobStatus")


Such parsing logic looks only limited to Sagemaker, right ? Does this PR only support Sagemaker ? Suggest add some comment

Its both sagemaker and OpenAI. Will add more comments

Zhangxunmt · 2024-09-03T20:40:18Z

common/src/main/java/org/opensearch/ml/common/MLTask.java

+        if (streamOutputVersion.onOrAfter(MLTask.MINIMAL_SUPPORTED_VERSION_FOR_BATCH_TRANSFORM_JOB)) {
+            if (remoteJob != null) {
+                out.writeBoolean(true);
+                try {
+                    AccessController.doPrivileged((PrivilegedExceptionAction<Void>) () -> {
+                        out.writeString(gson.toJson(remoteJob));
+                        return null;
+                    });
+                } catch (PrivilegedActionException e) {
+                    throw new RuntimeException(e);
+                }
+            } else {
+                out.writeBoolean(false);
+            }
+        }


You can use these two lines to replace this code block.
out.writeBoolean(true);
out.writeMap(remoteJob, StreamOutput::writeString, StreamOutput::writeGenericValue);

Zhangxunmt · 2024-09-03T20:41:11Z

common/src/main/java/org/opensearch/ml/common/MLTask.java

+            if (input.readBoolean()) {
+                String mapStr = input.readString();
+                this.remoteJob = gson.fromJson(mapStr, Map.class);
+            }


this.remoteJob = input.readMap(s -> s.readString(), s -> s.readGenericValue());
Why not use this for simplicity?

Zhangxunmt · 2024-09-03T20:50:21Z

plugin/src/main/java/org/opensearch/ml/action/tasks/CancelBatchJobTransportAction.java

+        Map<String, String> parameters = new HashMap<>();
+        for (Map.Entry<String, ?> entry : remoteJob.entrySet()) {
+            if (entry.getValue() instanceof String) {
+                parameters.put(entry.getKey(), (String) entry.getValue());
+            } else {
+                log.debug("Value for key " + entry.getKey() + " is not a String");
+            }
+        }


Seems like the remoteJob has to be <String, String>? If that's the case, why not throw exception when this remoteJob is created inside MLTask in the first place?

Yeah it seems easier when remote job is Map<String, String> while creating itself. Will make that change to avoid all this confusion

When converting to Map<String, String>, the values that are not strings will be lost. User may not be able to see complete response from remote job. I guess it is safe to use Map<String, Object> so that all key-value pairs from remote service are retained

Zhangxunmt · 2024-09-03T20:57:40Z

plugin/src/main/java/org/opensearch/ml/action/tasks/CancelBatchJobTransportAction.java

+        if (parameters.containsKey("TransformJobArn") && parameters.get("TransformJobArn") != null) {
+            String jobArn = parameters.get("TransformJobArn");
+            String transformJobName = jobArn.substring(jobArn.lastIndexOf("/") + 1);
+            parameters.put("TransformJobName", transformJobName);
+            parameters.remove("TransformJobArn");
+        }


parameters.put("TransformJobName", Optional.ofNullable(parameters.remove("TransformJobArn"))
.map(jobArn -> jobArn.substring(jobArn.lastIndexOf("/") + 1))
.orElse(null));

This block is used twice in the PR. Consider replacing it with this single statement to avoid code duplication.

Zhangxunmt · 2024-09-03T21:06:02Z

plugin/src/main/java/org/opensearch/ml/action/tasks/CancelBatchJobTransportAction.java

+                    actionListener.onResponse(new MLCancelBatchJobResponse(RestStatus.OK));
+                } else {
+                    log.debug("The status code from remote service is: " + modelOutput.getStatusCode());
+                    actionListener.onFailure(new ResourceNotFoundException("Couldn't cancel the transform job. Please try again"));


Why use ResourceNotFoundException here? It seems like that it can be any Exception from remote service. So we can just actionListener.onFailure(new OpenSearchException(modelOutput.getStatusCode()));

Yeah makes sense

Zhangxunmt · 2024-09-03T21:19:54Z

plugin/src/main/java/org/opensearch/ml/action/tasks/GetTaskTransportAction.java

+                        Map<String, Object> updatedTask = new HashMap<>();
+                        updatedTask.put(REMOTE_JOB_FIELD, remoteJob);
+
+                        if ((remoteJob.containsKey("status") && remoteJob.get("status").equals("completed"))


Is it possible that the key-value pair in remoteJobStatus from remote service is not "status" -> "completed" || "Completed", but some other names or status?

Yeah it can be other values like in progress, cancelled, etc

ylwu-amzn · 2024-09-05T00:56:01Z

* What went wrong:
Execution failed for task ':opensearch-ml-plugin:spotlessJavaCheck'.
> The following files had format violations:
      src/main/java/org/opensearch/ml/action/tasks/GetTaskTransportAction.java
          @@ -13,7 +13,6 @@
           import?static?org.opensearch.ml.common.MLTaskState.CANCELLED;
           import?static?org.opensearch.ml.common.MLTaskState.COMPLETED;
           import?static?org.opensearch.ml.common.connector.ConnectorAction.ActionType.BATCH_PREDICT_STATUS;
          -import?static?org.opensearch.ml.common.conversation.ConversationalIndexConstants.META_INDEX_NAME;
           import?static?org.opensearch.ml.utils.MLExceptionUtils.logException;
           import?static?org.opensearch.ml.utils.MLNodeUtils.createXContentParserFromRegistry;
           
          @@ -23,8 +22,6 @@
           
           import?org.opensearch.OpenSearchException;
           import?org.opensearch.OpenSearchStatusException;
          -import?org.opensearch.OpenSearchWrapperException;
          -import?org.opensearch.ResourceAlreadyExistsException;
           import?org.opensearch.ResourceNotFoundException;
           import?org.opensearch.action.ActionRequest;
           import?org.opensearch.action.get.GetRequest;
  Run './gradlew :opensearch-ml-plugin:spotlessApply' to fix these violations.

ylwu-amzn · 2024-09-05T02:07:28Z

REPRODUCE WITH: gradlew ':opensearch-ml-plugin:test' --tests "org.opensearch.ml.action.model_group.UpdateModelGroupITTests.test_update_public_model_group" -Dtests.seed=ACF849029908FA76 -Dtests.security.manager=false -Dtests.locale=ar-PS -Dtests.timezone=Asia/Taipei -Druntime.java=21
UpdateModelGroupITTests > test_update_public_model_group FAILED
    java.lang.AssertionError: 
    Expected: an instance of java.lang.IllegalArgumentException
         but: <org.junit.internal.runners.model.MultipleFailureException: There were ? errors:
      java.lang.IllegalArgumentException(You cannot specify model access control parameters because the Security plugin or model access control is disabled on your cluster.)

Signed-off-by: Bhavana Ramaram <[email protected]>

* support get batch transform job status in get task API Signed-off-by: Bhavana Ramaram <[email protected]> * add cancel batch prediction job API for offline inference Signed-off-by: Bhavana Ramaram <[email protected]> * add unit tests and address comments Signed-off-by: Bhavana Ramaram <[email protected]> * stash context for get model Signed-off-by: Bhavana Ramaram <[email protected]> * apply spotlessJava and exclude from test coverage Signed-off-by: Bhavana Ramaram <[email protected]> --------- Signed-off-by: Bhavana Ramaram <[email protected]> (cherry picked from commit 8da7bd2)

* support get batch transform job status in get task API Signed-off-by: Bhavana Ramaram <[email protected]> * add cancel batch prediction job API for offline inference Signed-off-by: Bhavana Ramaram <[email protected]> * add unit tests and address comments Signed-off-by: Bhavana Ramaram <[email protected]> * stash context for get model Signed-off-by: Bhavana Ramaram <[email protected]> * apply spotlessJava and exclude from test coverage Signed-off-by: Bhavana Ramaram <[email protected]> --------- Signed-off-by: Bhavana Ramaram <[email protected]> (cherry picked from commit 8da7bd2) Co-authored-by: Bhavana Ramaram <[email protected]>

rbhavna requested review from b4sjoo, dhrubo-os, jngz-es, model-collapse, ylwu-amzn, zane-neo, Zhangxunmt, austintlee, HenryL27, sam-herman and xinyual as code owners August 13, 2024 18:28

rbhavna had a problem deploying to ml-commons-cicd-env August 13, 2024 18:28 — with GitHub Actions Failure

rbhavna had a problem deploying to ml-commons-cicd-env August 13, 2024 18:29 — with GitHub Actions Failure

Zhangxunmt reviewed Aug 30, 2024

View reviewed changes

ylwu-amzn reviewed Aug 30, 2024

View reviewed changes

Zhangxunmt mentioned this pull request Aug 30, 2024

Batch inference #2687

Closed

5 tasks

Zhangxunmt added the 2.17 label Aug 30, 2024

ylwu-amzn reviewed Aug 30, 2024

View reviewed changes

rbhavna force-pushed the batch_task_status branch from 2e6186b to 2e616f0 Compare August 31, 2024 00:35

rbhavna had a problem deploying to ml-commons-cicd-env August 31, 2024 00:35 — with GitHub Actions Failure

rbhavna had a problem deploying to ml-commons-cicd-env August 31, 2024 00:36 — with GitHub Actions Failure

rbhavna force-pushed the batch_task_status branch from 2e616f0 to b6be1f5 Compare August 31, 2024 00:39

rbhavna had a problem deploying to ml-commons-cicd-env August 31, 2024 00:39 — with GitHub Actions Failure

Zhangxunmt reviewed Sep 3, 2024

View reviewed changes

Zhangxunmt previously approved these changes Sep 5, 2024

View reviewed changes

rbhavna dismissed Zhangxunmt’s stale review via a5e59a9 September 5, 2024 01:02

rbhavna had a problem deploying to ml-commons-cicd-env September 5, 2024 01:02 — with GitHub Actions Failure

rbhavna temporarily deployed to ml-commons-cicd-env September 5, 2024 01:03 — with GitHub Actions Inactive

rbhavna temporarily deployed to ml-commons-cicd-env September 5, 2024 02:15 — with GitHub Actions Inactive

rbhavna temporarily deployed to ml-commons-cicd-env September 5, 2024 14:31 — with GitHub Actions Inactive

rbhavna added 5 commits September 5, 2024 11:22

support get batch transform job status in get task API

edefb6a

Signed-off-by: Bhavana Ramaram <[email protected]>

add cancel batch prediction job API for offline inference

9a44d66

Signed-off-by: Bhavana Ramaram <[email protected]>

add unit tests and address comments

db7a334

Signed-off-by: Bhavana Ramaram <[email protected]>

stash context for get model

2e029a6

Signed-off-by: Bhavana Ramaram <[email protected]>

apply spotlessJava and exclude from test coverage

8460349

Signed-off-by: Bhavana Ramaram <[email protected]>

rbhavna force-pushed the batch_task_status branch from a5e59a9 to 8460349 Compare September 5, 2024 16:24

rbhavna temporarily deployed to ml-commons-cicd-env September 5, 2024 16:24 — with GitHub Actions Inactive

rbhavna had a problem deploying to ml-commons-cicd-env September 5, 2024 16:24 — with GitHub Actions Failure

ylwu-amzn approved these changes Sep 5, 2024

View reviewed changes

Zhangxunmt approved these changes Sep 5, 2024

View reviewed changes

Zhangxunmt merged commit 8da7bd2 into opensearch-project:main Sep 5, 2024
3 of 5 checks passed

Zhangxunmt added backport 2.x backport 2.17 labels Sep 5, 2024

This was referenced Sep 5, 2024

[Backport 2.x] support get batch transform job status in get task API #2893

Merged

[Backport 2.17] support get batch transform job status in get task API #2894

Merged

This was referenced Sep 6, 2024

[AUTOCUT] Integration Test failed for flow-framework: 2.17.0 opensearch-project/flow-framework#856

Closed

[BUG] Compile failure with new MLTask Constructor opensearch-project/flow-framework#858

Closed

junweid62 mentioned this pull request Sep 6, 2024

Fix and Refactor Test with Builder Pattern for Issue #2825 opensearch-project/flow-framework#859

Merged

2 tasks

		}

		private void processRemoteBatchPrediction(MLTask mlTask, String taskId, ActionListener<MLTaskGetResponse> actionListener) {

support get batch transform job status in get task API #2825

support get batch transform job status in get task API #2825

Conversation

rbhavna commented Aug 13, 2024 • edited Loading

Description

Related Issues

Check List

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Zhangxunmt commented Aug 30, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ylwu-amzn Aug 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ylwu-amzn commented Sep 5, 2024

ylwu-amzn commented Sep 5, 2024

rbhavna commented Aug 13, 2024 •

edited

Loading

ylwu-amzn Aug 30, 2024 •

edited

Loading