Change httpclient to async #1958
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -5,54 +5,43 @@ | |
|
|
||
| package org.opensearch.ml.engine.algorithms.remote; | ||
|
|
||
| import static org.opensearch.ml.common.connector.ConnectorProtocols.AWS_SIGV4; | ||
| import static software.amazon.awssdk.http.SdkHttpMethod.POST; | ||
|
|
||
| import java.security.AccessController; | ||
| import java.security.PrivilegedExceptionAction; | ||
| import java.time.Duration; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
| import java.util.concurrent.CompletableFuture; | ||
|
|
||
| import org.opensearch.client.Client; | ||
| import org.opensearch.common.util.TokenBucket; | ||
| import org.opensearch.core.action.ActionListener; | ||
| import org.opensearch.ml.common.connector.AwsConnector; | ||
| import org.opensearch.ml.common.connector.Connector; | ||
| import org.opensearch.ml.common.exception.MLException; | ||
| import org.opensearch.ml.common.input.MLInput; | ||
| import org.opensearch.ml.common.model.MLGuard; | ||
| import org.opensearch.ml.common.output.model.ModelTensors; | ||
| import org.opensearch.ml.engine.annotation.ConnectorExecutor; | ||
| import org.opensearch.ml.engine.httpclient.MLHttpClientFactory; | ||
| import org.opensearch.script.ScriptService; | ||
|
|
||
| import lombok.Getter; | ||
| import lombok.Setter; | ||
| import lombok.extern.log4j.Log4j2; | ||
| import software.amazon.awssdk.core.internal.http.async.SimpleHttpContentPublisher; | ||
| import software.amazon.awssdk.http.SdkHttpFullRequest; | ||
| import software.amazon.awssdk.http.async.AsyncExecuteRequest; | ||
| import software.amazon.awssdk.http.async.SdkAsyncHttpClient; | ||
|
|
||
| @Log4j2 | ||
| @ConnectorExecutor(AWS_SIGV4) | ||
| public class AwsConnectorExecutor extends AbstractConnectorExecutor { | ||
|
|
||
| @Getter | ||
| private AwsConnector connector; | ||
| @Setter | ||
| @Getter | ||
| private ScriptService scriptService; | ||
|
|
@@ -69,103 +58,52 @@ public class AwsConnectorExecutor extends AbstractConnectorExecutor { | |
| @Getter | ||
| private MLGuard mlGuard; | ||
|
|
||
| private SdkAsyncHttpClient httpClient; | ||
|
|
||
| public AwsConnectorExecutor(Connector connector) { | ||
| super.initialize(connector); | ||
| this.connector = (AwsConnector) connector; | ||
| Duration connectionTimeout = Duration.ofSeconds(super.getConnectorClientConfig().getConnectionTimeout()); | ||
| Duration readTimeout = Duration.ofSeconds(super.getConnectorClientConfig().getReadTimeout()); | ||
| Integer maxConnection = super.getConnectorClientConfig().getMaxConnections(); | ||
| this.httpClient = MLHttpClientFactory.getAsyncHttpClient(connectionTimeout, readTimeout, maxConnection); | ||
| } | ||
|
|
||
| @SuppressWarnings("removal") | ||
| @Override | ||
| public void invokeRemoteModel( | ||
| MLInput mlInput, | ||
| Map<String, String> parameters, | ||
| String payload, | ||
| Map<Integer, ModelTensors> tensorOutputs, | ||
| ExecutionContext countDownLatch, | ||
| ActionListener<List<ModelTensors>> actionListener | ||
| ) { | ||
| try { | ||
| SdkHttpFullRequest request = ConnectorUtils.buildSdkRequest(connector, parameters, payload, POST); | ||
|
There was a problem hiding this comment. This assumes all invoke remote model is a "Post" request by hardcoding. Very soon we will add more action types with different methods. So do you want to extend this to support more HTTP Method through different action types? There was a problem hiding this comment. This change is based on the existing code: https://github.com/opensearch-project/ml-commons/blob/main/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutor.java#L116, I think when we needs to support more different methods, we can change this at that time, current implementation can adapt to that change easily. |
||
| AsyncExecuteRequest executeRequest = AsyncExecuteRequest | ||
| .builder() | ||
| .request(signRequest(request)) | ||
| .requestContentPublisher(new SimpleHttpContentPublisher(request)) | ||
| .responseHandler( | ||
| new MLSdkAsyncHttpResponseHandler( | ||
| countDownLatch, | ||
| actionListener, | ||
| parameters, | ||
| tensorOutputs, | ||
| connector, | ||
| scriptService, | ||
| mlGuard | ||
| ) | ||
| ) | ||
| .build(); | ||
| AccessController.doPrivileged((PrivilegedExceptionAction<CompletableFuture<Void>>) () -> httpClient.execute(executeRequest)); | ||
| } catch (RuntimeException exception) { | ||
| log.error("Failed to execute predict in aws connector: " + exception.getMessage(), exception); | ||
| actionListener.onFailure(exception); | ||
| } catch (Throwable e) { | ||
| log.error("Failed to execute predict in aws connector", e); | ||
| actionListener.onFailure(new MLException("Fail to execute predict in aws connector", e)); | ||
| } | ||
| } | ||
|
|
||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,32 @@ | ||
| /* | ||
| * | ||
| * * Copyright OpenSearch Contributors | ||
| * * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| */ | ||
|
|
||
| package org.opensearch.ml.engine.algorithms.remote; | ||
|
|
||
| import java.util.concurrent.CountDownLatch; | ||
| import java.util.concurrent.atomic.AtomicReference; | ||
|
|
||
| import lombok.AllArgsConstructor; | ||
| import lombok.Data; | ||
|
|
||
| /** | ||
| * This class encapsulates several parameters that are used in a split-batch request case. | ||
| * A batch request is that in neural-search side multiple fields are send in one request to ml-commons, | ||
| * but the remote model doesn't accept list of string inputs so in ml-commons the request needs split. | ||
| * sequence is used to identify the index of the split request. | ||
| * countDownLatch is used to wait for all the split requests to finish. | ||
| * exceptionHolder is used to hold any exception thrown in a split-batch request. | ||
| */ | ||
| @Data | ||
| @AllArgsConstructor | ||
| public class ExecutionContext { | ||
|
There was a problem hiding this comment. Let's add some comment about what is this class for? |
||
| // Should never be null | ||
| private int sequence; | ||
| private CountDownLatch countDownLatch; | ||
| // This is to hold any exception thrown in a split-batch request | ||
| private AtomicReference<Exception> exceptionHolder; | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Codecov shows this default exception is not covered. Also what is the necessity here to default an exception? When the method is not implemented, the compile would fail first?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Added tests for this. No, this is a runtime error, compilation won't be impacted by this. Below are reasons:
First of all, this is an interface method but not all subcalsses need to implement this method, e.g. RemoteModel. So we need to make it a default method, and this method can be accessed in all its implemented classes. Also this method needs return value, there are three options:
In fact, there's another option which is keep this method abstract and all subclasses needs to implement this and based on their case, they can throw UnsupportedOperationException which can also avoid accidentally abuse of the method, but this option we need to change quite a lot of files, and since this also works I choose this simpler option.