Skip to content

Commit 2831130

Browse files
brianf-awsbrianf-aws
committed
apply spotles
Signed-off-by: Brian Flores <[email protected]>
1 parent 8a69617 commit 2831130

File tree

2 files changed

+58
-54
lines changed

2 files changed

+58
-54
lines changed

‎plugin/src/main/java/org/opensearch/ml/action/undeploy/TransportUndeployModelsAction.java‎

Lines changed: 30 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -222,10 +222,10 @@ private void undeployModels(
222222
}
223223

224224
private void bulkSetModelIndexToUndeploy(
225-
String[] modelIds,
226-
String tenantId,
227-
ActionListener<MLUndeployModelsResponse> listener,
228-
MLUndeployModelNodesResponse mlUndeployModelNodesResponse
225+
String[] modelIds,
226+
String tenantId,
227+
ActionListener<MLUndeployModelsResponse> listener,
228+
MLUndeployModelNodesResponse mlUndeployModelNodesResponse
229229
) {
230230
BulkDataObjectRequest bulkRequest = BulkDataObjectRequest.builder().globalIndex(ML_MODEL_INDEX).build();
231231

@@ -240,11 +240,11 @@ private void bulkSetModelIndexToUndeploy(
240240
updateDocument.put(MLModel.CURRENT_WORKER_NODE_COUNT_FIELD, 0);
241241

242242
UpdateDataObjectRequest updateRequest = UpdateDataObjectRequest
243-
.builder()
244-
.id(modelId)
245-
.tenantId(tenantId)
246-
.dataObject(updateDocument)
247-
.build();
243+
.builder()
244+
.id(modelId)
245+
.tenantId(tenantId)
246+
.dataObject(updateDocument)
247+
.build();
248248
bulkRequest.add(updateRequest).setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
249249
}
250250

@@ -258,12 +258,12 @@ private void bulkSetModelIndexToUndeploy(
258258
listenerWithContextRestoration.onResponse(new MLUndeployModelsResponse(mlUndeployModelNodesResponse));
259259
}, e -> {
260260
String modelsNotFoundMessage = String
261-
.format("Failed to set the following modelId(s) to UNDEPLOY in index: %s", Arrays.toString(modelIds));
261+
.format("Failed to set the following modelId(s) to UNDEPLOY in index: %s", Arrays.toString(modelIds));
262262
log.error(modelsNotFoundMessage, e);
263263

264264
OpenSearchStatusException exception = new OpenSearchStatusException(
265-
modelsNotFoundMessage + e.getMessage(),
266-
RestStatus.INTERNAL_SERVER_ERROR
265+
modelsNotFoundMessage + e.getMessage(),
266+
RestStatus.INTERNAL_SERVER_ERROR
267267
);
268268
listenerWithContextRestoration.onFailure(exception);
269269
});
@@ -278,19 +278,24 @@ private void bulkSetModelIndexToUndeploy(
278278
try {
279279
BulkResponse bulkResponse = BulkResponse.fromXContent(response.parser());
280280
log
281-
.info(
282-
"Executed {} bulk operations with {} failures, Took: {}",
283-
bulkResponse.getItems().length,
284-
bulkResponse.hasFailures()
285-
? Arrays.stream(bulkResponse.getItems()).filter(BulkItemResponse::isFailed).count()
286-
: 0,
287-
bulkResponse.getTook()
288-
);
289-
List<String> unemployedModelIds = Arrays.stream(bulkResponse.getItems())
290-
.filter(bulkItemResponse -> !bulkItemResponse.isFailed())
291-
.map(BulkItemResponse::getId)
292-
.collect(Collectors.toList());
293-
log.debug("Successfully set the following modelId(s) to UNDEPLOY in index: {}", Arrays.toString(unemployedModelIds.toArray()));
281+
.info(
282+
"Executed {} bulk operations with {} failures, Took: {}",
283+
bulkResponse.getItems().length,
284+
bulkResponse.hasFailures()
285+
? Arrays.stream(bulkResponse.getItems()).filter(BulkItemResponse::isFailed).count()
286+
: 0,
287+
bulkResponse.getTook()
288+
);
289+
List<String> unemployedModelIds = Arrays
290+
.stream(bulkResponse.getItems())
291+
.filter(bulkItemResponse -> !bulkItemResponse.isFailed())
292+
.map(BulkItemResponse::getId)
293+
.collect(Collectors.toList());
294+
log
295+
.debug(
296+
"Successfully set the following modelId(s) to UNDEPLOY in index: {}",
297+
Arrays.toString(unemployedModelIds.toArray())
298+
);
294299

295300
bulkResponseListener.onResponse(bulkResponse);
296301
} catch (Exception e) {

‎plugin/src/test/java/org/opensearch/ml/action/undeploy/TransportUndeployModelsActionTests.java‎

Lines changed: 28 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,27 @@
77

88
package org.opensearch.ml.action.undeploy;
99

10+
import static org.mockito.ArgumentMatchers.any;
11+
import static org.mockito.ArgumentMatchers.isA;
12+
import static org.mockito.Mockito.doAnswer;
13+
import static org.mockito.Mockito.doReturn;
14+
import static org.mockito.Mockito.doThrow;
15+
import static org.mockito.Mockito.mock;
16+
import static org.mockito.Mockito.never;
17+
import static org.mockito.Mockito.spy;
18+
import static org.mockito.Mockito.verify;
19+
import static org.mockito.Mockito.when;
20+
import static org.opensearch.ml.common.CommonValue.ML_MODEL_INDEX;
21+
import static org.opensearch.ml.common.CommonValue.NOT_FOUND;
22+
import static org.opensearch.ml.task.MLPredictTaskRunnerTests.USER_STRING;
23+
24+
import java.io.IOException;
25+
import java.util.ArrayList;
26+
import java.util.Collections;
27+
import java.util.HashMap;
28+
import java.util.List;
29+
import java.util.Map;
30+
1031
import org.junit.Before;
1132
import org.junit.Rule;
1233
import org.junit.rules.ExpectedException;
@@ -54,27 +75,6 @@
5475
import org.opensearch.threadpool.ThreadPool;
5576
import org.opensearch.transport.TransportService;
5677

57-
import java.io.IOException;
58-
import java.util.ArrayList;
59-
import java.util.Collections;
60-
import java.util.HashMap;
61-
import java.util.List;
62-
import java.util.Map;
63-
64-
import static org.mockito.ArgumentMatchers.any;
65-
import static org.mockito.ArgumentMatchers.isA;
66-
import static org.mockito.Mockito.doAnswer;
67-
import static org.mockito.Mockito.doReturn;
68-
import static org.mockito.Mockito.doThrow;
69-
import static org.mockito.Mockito.mock;
70-
import static org.mockito.Mockito.never;
71-
import static org.mockito.Mockito.spy;
72-
import static org.mockito.Mockito.verify;
73-
import static org.mockito.Mockito.when;
74-
import static org.opensearch.ml.common.CommonValue.ML_MODEL_INDEX;
75-
import static org.opensearch.ml.common.CommonValue.NOT_FOUND;
76-
import static org.opensearch.ml.task.MLPredictTaskRunnerTests.USER_STRING;
77-
7878
public class TransportUndeployModelsActionTests extends OpenSearchTestCase {
7979

8080
@Mock
@@ -547,14 +547,13 @@ public void testDoExecute_modelIds_moreThan1() {
547547

548548
private BulkResponse getSuccessBulkResponse() {
549549
return new BulkResponse(
550-
new BulkItemResponse[]{
551-
new BulkItemResponse(
552-
1,
553-
DocWriteRequest.OpType.UPDATE,
554-
new UpdateResponse(new ShardId(ML_MODEL_INDEX, "modelId123", 0), "id1", 1, 1, 1, DocWriteResponse.Result.UPDATED)
555-
)
556-
},
557-
100L
550+
new BulkItemResponse[] {
551+
new BulkItemResponse(
552+
1,
553+
DocWriteRequest.OpType.UPDATE,
554+
new UpdateResponse(new ShardId(ML_MODEL_INDEX, "modelId123", 0), "id1", 1, 1, 1, DocWriteResponse.Result.UPDATED)
555+
) },
556+
100L
558557
);
559558
}
560559
}

0 commit comments

Comments
 (0)