Skip to content

Commit

Permalink
fix get batch task bug (#2937)
Browse files Browse the repository at this point in the history
Signed-off-by: Yaliang Wu <ylwu@amazon.com>
  • Loading branch information
ylwu-amzn authored Sep 12, 2024
1 parent 30228ea commit bbaea1c
Showing 1 changed file with 14 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,13 @@
package org.opensearch.ml.engine.algorithms.remote;

import static org.opensearch.ml.common.connector.ConnectorProtocols.AWS_SIGV4;
import static software.amazon.awssdk.http.SdkHttpMethod.GET;
import static software.amazon.awssdk.http.SdkHttpMethod.POST;

import java.security.AccessController;
import java.security.PrivilegedExceptionAction;
import java.time.Duration;
import java.util.Locale;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

Expand Down Expand Up @@ -86,7 +88,18 @@ public void invokeRemoteService(
ActionListener<Tuple<Integer, ModelTensors>> actionListener
) {
try {
SdkHttpFullRequest request = ConnectorUtils.buildSdkRequest(action, connector, parameters, payload, POST);
SdkHttpFullRequest request;
switch (connector.getActionHttpMethod(action).toUpperCase(Locale.ROOT)) {
case "POST":
log.debug("original payload to remote model: " + payload);
request = ConnectorUtils.buildSdkRequest(action, connector, parameters, payload, POST);
break;
case "GET":
request = ConnectorUtils.buildSdkRequest(action, connector, parameters, null, GET);
break;
default:
throw new IllegalArgumentException("unsupported http method");
}
AsyncExecuteRequest executeRequest = AsyncExecuteRequest
.builder()
.request(signRequest(request))
Expand Down

0 comments on commit bbaea1c

Please sign in to comment.