Skip to content

Commit c2c9287

Browse files
brianf-awsbrianf-aws
committed
apply spotless
Signed-off-by: Brian Flores <[email protected]>
1 parent 3ee2723 commit c2c9287

File tree

2 files changed

+61
-56
lines changed

2 files changed

+61
-56
lines changed

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

Lines changed: 32 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import org.opensearch.action.support.ActionFilters;
2626
import org.opensearch.action.support.HandledTransportAction;
2727
import org.opensearch.action.support.WriteRequest;
28-
import org.opensearch.client.Client;
2928
import org.opensearch.cluster.service.ClusterService;
3029
import org.opensearch.common.inject.Inject;
3130
import org.opensearch.common.settings.Settings;
@@ -66,6 +65,7 @@
6665
import org.opensearch.tasks.Task;
6766
import org.opensearch.threadpool.ThreadPool;
6867
import org.opensearch.transport.TransportService;
68+
import org.opensearch.transport.client.Client;
6969

7070
import com.google.common.annotations.VisibleForTesting;
7171

@@ -218,6 +218,7 @@ private void undeployModels(
218218
return modelCacheMissForModelIds;
219219
});
220220
if (response.getNodes().isEmpty() || modelNotFoundInNodesCache) {
221+
log.warn("No nodes service these models, performing manual `UNDEPLOY` write to model index");
221222
bulkSetModelIndexToUndeploy(modelIds, tenantId, listener, response);
222223
return;
223224
}
@@ -227,10 +228,10 @@ private void undeployModels(
227228
}
228229

229230
private void bulkSetModelIndexToUndeploy(
230-
String[] modelIds,
231-
String tenantId,
232-
ActionListener<MLUndeployModelsResponse> listener,
233-
MLUndeployModelNodesResponse mlUndeployModelNodesResponse
231+
String[] modelIds,
232+
String tenantId,
233+
ActionListener<MLUndeployModelsResponse> listener,
234+
MLUndeployModelNodesResponse mlUndeployModelNodesResponse
234235
) {
235236
BulkDataObjectRequest bulkRequest = BulkDataObjectRequest.builder().globalIndex(ML_MODEL_INDEX).build();
236237

@@ -245,11 +246,11 @@ private void bulkSetModelIndexToUndeploy(
245246
updateDocument.put(MLModel.CURRENT_WORKER_NODE_COUNT_FIELD, 0);
246247

247248
UpdateDataObjectRequest updateRequest = UpdateDataObjectRequest
248-
.builder()
249-
.id(modelId)
250-
.tenantId(tenantId)
251-
.dataObject(updateDocument)
252-
.build();
249+
.builder()
250+
.id(modelId)
251+
.tenantId(tenantId)
252+
.dataObject(updateDocument)
253+
.build();
253254
bulkRequest.add(updateRequest).setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
254255
}
255256

@@ -263,12 +264,12 @@ private void bulkSetModelIndexToUndeploy(
263264
listenerWithContextRestoration.onResponse(new MLUndeployModelsResponse(mlUndeployModelNodesResponse));
264265
}, e -> {
265266
String modelsNotFoundMessage = String
266-
.format("Failed to set the following modelId(s) to UNDEPLOY in index: %s", Arrays.toString(modelIds));
267+
.format("Failed to set the following modelId(s) to UNDEPLOY in index: %s", Arrays.toString(modelIds));
267268
log.error(modelsNotFoundMessage, e);
268269

269270
OpenSearchStatusException exception = new OpenSearchStatusException(
270-
modelsNotFoundMessage + e.getMessage(),
271-
RestStatus.INTERNAL_SERVER_ERROR
271+
modelsNotFoundMessage + e.getMessage(),
272+
RestStatus.INTERNAL_SERVER_ERROR
272273
);
273274
listenerWithContextRestoration.onFailure(exception);
274275
});
@@ -283,19 +284,24 @@ private void bulkSetModelIndexToUndeploy(
283284
try {
284285
BulkResponse bulkResponse = BulkResponse.fromXContent(response.parser());
285286
log
286-
.info(
287-
"Executed {} bulk operations with {} failures, Took: {}",
288-
bulkResponse.getItems().length,
289-
bulkResponse.hasFailures()
290-
? Arrays.stream(bulkResponse.getItems()).filter(BulkItemResponse::isFailed).count()
291-
: 0,
292-
bulkResponse.getTook()
293-
);
294-
List<String> unemployedModelIds = Arrays.stream(bulkResponse.getItems())
295-
.filter(bulkItemResponse -> !bulkItemResponse.isFailed())
296-
.map(BulkItemResponse::getId)
297-
.collect(Collectors.toList());
298-
log.debug("Successfully set the following modelId(s) to UNDEPLOY in index: {}", Arrays.toString(unemployedModelIds.toArray()));
287+
.info(
288+
"Executed {} bulk operations with {} failures, Took: {}",
289+
bulkResponse.getItems().length,
290+
bulkResponse.hasFailures()
291+
? Arrays.stream(bulkResponse.getItems()).filter(BulkItemResponse::isFailed).count()
292+
: 0,
293+
bulkResponse.getTook()
294+
);
295+
List<String> unemployedModelIds = Arrays
296+
.stream(bulkResponse.getItems())
297+
.filter(bulkItemResponse -> !bulkItemResponse.isFailed())
298+
.map(BulkItemResponse::getId)
299+
.collect(Collectors.toList());
300+
log
301+
.debug(
302+
"Successfully set the following modelId(s) to UNDEPLOY in index: {}",
303+
Arrays.toString(unemployedModelIds.toArray())
304+
);
299305

300306
bulkResponseListener.onResponse(bulkResponse);
301307
} catch (Exception e) {

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

Lines changed: 29 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,27 @@
55

66
package org.opensearch.ml.action.undeploy;
77

8+
import static org.mockito.ArgumentMatchers.any;
9+
import static org.mockito.ArgumentMatchers.isA;
10+
import static org.mockito.Mockito.doAnswer;
11+
import static org.mockito.Mockito.doReturn;
12+
import static org.mockito.Mockito.doThrow;
13+
import static org.mockito.Mockito.mock;
14+
import static org.mockito.Mockito.never;
15+
import static org.mockito.Mockito.spy;
16+
import static org.mockito.Mockito.verify;
17+
import static org.mockito.Mockito.when;
18+
import static org.opensearch.ml.common.CommonValue.ML_MODEL_INDEX;
19+
import static org.opensearch.ml.common.CommonValue.NOT_FOUND;
20+
import static org.opensearch.ml.task.MLPredictTaskRunnerTests.USER_STRING;
21+
22+
import java.io.IOException;
23+
import java.util.ArrayList;
24+
import java.util.Collections;
25+
import java.util.HashMap;
26+
import java.util.List;
27+
import java.util.Map;
28+
829
import org.junit.Before;
930
import org.junit.Rule;
1031
import org.junit.rules.ExpectedException;
@@ -21,7 +42,6 @@
2142
import org.opensearch.action.support.ActionFilters;
2243
import org.opensearch.action.update.UpdateRequest;
2344
import org.opensearch.action.update.UpdateResponse;
24-
import org.opensearch.client.Client;
2545
import org.opensearch.cluster.ClusterName;
2646
import org.opensearch.cluster.service.ClusterService;
2747
import org.opensearch.common.settings.Settings;
@@ -51,27 +71,7 @@
5171
import org.opensearch.test.OpenSearchTestCase;
5272
import org.opensearch.threadpool.ThreadPool;
5373
import org.opensearch.transport.TransportService;
54-
55-
import java.io.IOException;
56-
import java.util.ArrayList;
57-
import java.util.Collections;
58-
import java.util.HashMap;
59-
import java.util.List;
60-
import java.util.Map;
61-
62-
import static org.mockito.ArgumentMatchers.any;
63-
import static org.mockito.ArgumentMatchers.isA;
64-
import static org.mockito.Mockito.doAnswer;
65-
import static org.mockito.Mockito.doReturn;
66-
import static org.mockito.Mockito.doThrow;
67-
import static org.mockito.Mockito.mock;
68-
import static org.mockito.Mockito.never;
69-
import static org.mockito.Mockito.spy;
70-
import static org.mockito.Mockito.verify;
71-
import static org.mockito.Mockito.when;
72-
import static org.opensearch.ml.common.CommonValue.ML_MODEL_INDEX;
73-
import static org.opensearch.ml.common.CommonValue.NOT_FOUND;
74-
import static org.opensearch.ml.task.MLPredictTaskRunnerTests.USER_STRING;
74+
import org.opensearch.transport.client.Client;
7575

7676
public class TransportUndeployModelsActionTests extends OpenSearchTestCase {
7777

@@ -545,14 +545,13 @@ public void testDoExecute_modelIds_moreThan1() {
545545

546546
private BulkResponse getSuccessBulkResponse() {
547547
return new BulkResponse(
548-
new BulkItemResponse[]{
549-
new BulkItemResponse(
550-
1,
551-
DocWriteRequest.OpType.UPDATE,
552-
new UpdateResponse(new ShardId(ML_MODEL_INDEX, "modelId123", 0), "id1", 1, 1, 1, DocWriteResponse.Result.UPDATED)
553-
)
554-
},
555-
100L
548+
new BulkItemResponse[] {
549+
new BulkItemResponse(
550+
1,
551+
DocWriteRequest.OpType.UPDATE,
552+
new UpdateResponse(new ShardId(ML_MODEL_INDEX, "modelId123", 0), "id1", 1, 1, 1, DocWriteResponse.Result.UPDATED)
553+
) },
554+
100L
556555
);
557556
}
558557
}

0 commit comments

Comments
 (0)