Skip to content

Commit f71950e

Browse files
b4sjoob4sjoo
committedMay 27, 2024
Add retry to delete model API
Signed-off-by: Sicheng Song <sicheng.song@outlook.com>
1 parent bb13148 commit f71950e

File tree

7 files changed

+668
-156
lines changed

7 files changed

+668
-156
lines changed
 

‎common/src/main/java/org/opensearch/ml/common/transport/model/MLModelDeleteRequest.java

+31-3
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import lombok.Getter;
1010
import org.opensearch.action.ActionRequest;
1111
import org.opensearch.action.ActionRequestValidationException;
12+
import org.opensearch.common.unit.TimeValue;
1213
import org.opensearch.core.common.io.stream.InputStreamStreamInput;
1314
import org.opensearch.core.common.io.stream.OutputStreamStreamOutput;
1415
import org.opensearch.core.common.io.stream.StreamInput;
@@ -21,24 +22,40 @@
2122

2223
import static org.opensearch.action.ValidateActions.addValidationError;
2324

25+
@Getter
2426
public class MLModelDeleteRequest extends ActionRequest {
25-
@Getter
2627
String modelId;
28+
Boolean retry;
29+
Integer maxRetry;
30+
TimeValue retryDelay;
31+
TimeValue retryTimeout;
2732

2833
@Builder
29-
public MLModelDeleteRequest(String modelId) {
34+
public MLModelDeleteRequest(String modelId, Boolean retry, Integer maxRetry, TimeValue retryDelay, TimeValue retryTimeout) {
3035
this.modelId = modelId;
36+
this.retry = retry;
37+
this.maxRetry = maxRetry;
38+
this.retryDelay = retryDelay;
39+
this.retryTimeout = retryTimeout;
3140
}
3241

3342
public MLModelDeleteRequest(StreamInput input) throws IOException {
3443
super(input);
3544
this.modelId = input.readString();
45+
this.retry = input.readOptionalBoolean();
46+
this.maxRetry = input.readOptionalVInt();
47+
this.retryDelay = input.readOptionalTimeValue();
48+
this.retryTimeout = input.readOptionalTimeValue();
3649
}
3750

3851
@Override
3952
public void writeTo(StreamOutput output) throws IOException {
4053
super.writeTo(output);
4154
output.writeString(modelId);
55+
output.writeOptionalBoolean(retry);
56+
output.writeOptionalVInt(maxRetry);
57+
output.writeOptionalTimeValue(retryDelay);
58+
output.writeOptionalTimeValue(retryTimeout);
4259
}
4360

4461
@Override
@@ -48,6 +65,17 @@ public ActionRequestValidationException validate() {
4865
if (this.modelId == null) {
4966
exception = addValidationError("ML model id can't be null", exception);
5067
}
68+
if (this.maxRetry != null && (this.maxRetry < 0 || this.maxRetry > 5)) {
69+
exception = addValidationError("Retry count should be greater than or equal to 0 and less than 5", exception);
70+
}
71+
72+
if (this.retryDelay != null && (this.retryDelay.getMillis() < 0 || this.retryDelay.getMillis() > 30000)) {
73+
exception = addValidationError("Retry delay should be greater than or equal to 0 or less than 30000 milliseconds", exception);
74+
}
75+
76+
if (this.retryTimeout != null && (this.retryTimeout.getMillis() < 0 || this.retryTimeout.getMillis() > 30000)) {
77+
exception = addValidationError("Retry delay should be greater than or equal to 0 or less than 30000 milliseconds", exception);
78+
}
5179

5280
return exception;
5381
}
@@ -58,7 +86,7 @@ public static MLModelDeleteRequest fromActionRequest(ActionRequest actionRequest
5886
}
5987

6088
try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
61-
OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
89+
OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) {
6290
actionRequest.writeTo(osso);
6391
try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) {
6492
return new MLModelDeleteRequest(input);

‎common/src/test/java/org/opensearch/ml/common/transport/model/MLModelDeleteRequestTest.java

+58
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import org.opensearch.action.ActionRequest;
1111
import org.opensearch.action.ActionRequestValidationException;
1212
import org.opensearch.common.io.stream.BytesStreamOutput;
13+
import org.opensearch.common.unit.TimeValue;
1314
import org.opensearch.core.common.io.stream.StreamOutput;
1415

1516
import java.io.IOException;
@@ -19,6 +20,7 @@
1920
import static org.junit.Assert.assertNotSame;
2021
import static org.junit.Assert.assertNull;
2122
import static org.junit.Assert.assertSame;
23+
import static org.junit.Assert.assertTrue;
2224

2325
public class MLModelDeleteRequestTest {
2426
private String modelId;
@@ -54,6 +56,62 @@ public void validate_Exception_NullModelId() {
5456
assertEquals("Validation Failed: 1: ML model id can't be null;", exception.getMessage());
5557
}
5658

59+
@Test
60+
public void validate_Exception_NegativeMaxRetry() {
61+
MLModelDeleteRequest mlModelDeleteRequest = MLModelDeleteRequest.builder()
62+
.modelId(modelId).maxRetry(-1).build();
63+
64+
ActionRequestValidationException exception = mlModelDeleteRequest.validate();
65+
assertEquals("Validation Failed: 1: Retry count should be greater than or equal to 0 and less than 5;", exception.getMessage());
66+
}
67+
68+
@Test
69+
public void validate_Exception_ExceedMaxRetry() {
70+
MLModelDeleteRequest mlModelDeleteRequest = MLModelDeleteRequest.builder()
71+
.modelId(modelId).maxRetry(6).build();
72+
73+
ActionRequestValidationException exception = mlModelDeleteRequest.validate();
74+
assertEquals("Validation Failed: 1: Retry count should be greater than or equal to 0 and less than 5;", exception.getMessage());
75+
}
76+
77+
@Test
78+
public void validate_Exception_NegativeRetryDelay() {
79+
MLModelDeleteRequest mlModelDeleteRequest = MLModelDeleteRequest.builder()
80+
.modelId(modelId).retryDelay(TimeValue.timeValueMillis(-1)).build();
81+
82+
ActionRequestValidationException exception = mlModelDeleteRequest.validate();
83+
assertEquals("Validation Failed: 1: Retry delay should be greater than or equal to 0 or less than 30000 milliseconds;", exception.getMessage());
84+
}
85+
86+
87+
@Test
88+
public void validate_Exception_ExceedRetryDelay() {
89+
MLModelDeleteRequest mlModelDeleteRequest = MLModelDeleteRequest.builder()
90+
.modelId(modelId).retryDelay(TimeValue.timeValueMillis(50000)).build();
91+
92+
ActionRequestValidationException exception = mlModelDeleteRequest.validate();
93+
assertEquals("Validation Failed: 1: Retry delay should be greater than or equal to 0 or less than 30000 milliseconds;", exception.getMessage());
94+
}
95+
96+
@Test
97+
public void validate_Exception_NegativeRetryTimeout() {
98+
MLModelDeleteRequest mlModelDeleteRequest = MLModelDeleteRequest.builder()
99+
.modelId(modelId).retryTimeout(TimeValue.timeValueMillis(-1)).build();
100+
101+
ActionRequestValidationException exception = mlModelDeleteRequest.validate();
102+
assertEquals("Validation Failed: 1: Retry delay should be greater than or equal to 0 or less than 30000 milliseconds;", exception.getMessage());
103+
}
104+
105+
106+
@Test
107+
public void validate_Exception_ExceedRetryTimeout() {
108+
MLModelDeleteRequest mlModelDeleteRequest = MLModelDeleteRequest.builder()
109+
.modelId(modelId).retryTimeout(TimeValue.timeValueSeconds(60)).build();
110+
111+
ActionRequestValidationException exception = mlModelDeleteRequest.validate();
112+
assertEquals("Validation Failed: 1: Retry delay should be greater than or equal to 0 or less than 30000 milliseconds;", exception.getMessage());
113+
}
114+
57115
@Test
58116
public void fromActionRequest_Success() {
59117
MLModelDeleteRequest mlModelDeleteRequest = MLModelDeleteRequest.builder()

0 commit comments

Comments
 (0)