Skip to content

Commit 2049b92

Browse files
committed
update error code to 429 for rate limiting and update logs
Signed-off-by: Xun Zhang <xunzh@amazon.com>
1 parent a22e926 commit 2049b92

File tree

3 files changed

+11
-6
lines changed

3 files changed

+11
-6
lines changed

plugin/src/main/java/org/opensearch/ml/action/batch/TransportBatchIngestionAction.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,6 @@
3333
import org.opensearch.ml.common.MLTask;
3434
import org.opensearch.ml.common.MLTaskState;
3535
import org.opensearch.ml.common.MLTaskType;
36-
import org.opensearch.ml.common.exception.MLLimitExceededException;
3736
import org.opensearch.ml.common.transport.batch.MLBatchIngestionAction;
3837
import org.opensearch.ml.common.transport.batch.MLBatchIngestionInput;
3938
import org.opensearch.ml.common.transport.batch.MLBatchIngestionRequest;
@@ -146,9 +145,10 @@ protected void createMLTaskandExecute(MLBatchIngestionInput mlBatchIngestionInpu
146145

147146
mlModelManager.checkMaxBatchJobTask(mlTask, ActionListener.wrap(exceedLimits -> {
148147
if (exceedLimits) {
149-
String error = "exceed maximum BATCH_INGEST Task limits";
148+
String error =
149+
"Exceeded maximum limit for BATCH_INGEST tasks. To increase the limit, update the plugins.ml_commons.max_batch_ingestion_tasks setting.";
150150
log.warn(error + " in task " + mlTask.getTaskId());
151-
listener.onFailure(new MLLimitExceededException(error));
151+
listener.onFailure(new OpenSearchStatusException(error, RestStatus.TOO_MANY_REQUESTS));
152152
} else {
153153
mlTaskManager.createMLTask(mlTask, ActionListener.wrap(response -> {
154154
String taskId = response.getId();

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

+3-3
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,6 @@
5555
import org.opensearch.ml.common.dataset.MLInputDataType;
5656
import org.opensearch.ml.common.dataset.MLInputDataset;
5757
import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
58-
import org.opensearch.ml.common.exception.MLLimitExceededException;
5958
import org.opensearch.ml.common.input.MLInput;
6059
import org.opensearch.ml.common.output.MLOutput;
6160
import org.opensearch.ml.common.output.MLPredictionOutput;
@@ -257,9 +256,10 @@ protected void executeTask(MLPredictionTaskRequest request, ActionListener<MLTas
257256
if (actionType.equals(ActionType.BATCH_PREDICT)) {
258257
mlModelManager.checkMaxBatchJobTask(mlTask, ActionListener.wrap(exceedLimits -> {
259258
if (exceedLimits) {
260-
String error = "exceed maximum BATCH_PREDICTION Task limits";
259+
String error =
260+
"Exceeded maximum limit for BATCH_PREDICTION tasks. To increase the limit, update the plugins.ml_commons.max_batch_inference_tasks setting.";
261261
log.warn(error + " in task " + mlTask.getTaskId());
262-
listener.onFailure(new MLLimitExceededException(error));
262+
listener.onFailure(new OpenSearchStatusException(error, RestStatus.TOO_MANY_REQUESTS));
263263
} else {
264264
executePredictionByInputDataType(inputDataType, modelId, mlInput, mlTask, functionName, listener);
265265
}

plugin/src/test/java/org/opensearch/ml/task/MLPredictTaskRunnerTests.java

+5
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,11 @@ public void setup() throws IOException {
229229

230230
GetResult getResult = new GetResult(indexName, "1.1.1", 111l, 111l, 111l, true, bytesReference, null, null);
231231
getResponse = new GetResponse(getResult);
232+
doAnswer(invocation -> {
233+
ActionListener<Boolean> listener = invocation.getArgument(1);
234+
listener.onResponse(false);
235+
return null;
236+
}).when(mlModelManager).checkMaxBatchJobTask(any(MLTask.class), isA(ActionListener.class));
232237
}
233238

234239
public void testExecuteTask_OnLocalNode() {

0 commit comments

Comments
 (0)