|
@@ -7,8 +7,12 @@
|
|
|
|
|
|
package org.elasticsearch.xpack.inference.action;
|
|
|
|
|
|
+import org.apache.logging.log4j.LogManager;
|
|
|
+import org.apache.logging.log4j.Logger;
|
|
|
+import org.elasticsearch.ElasticsearchStatusException;
|
|
|
import org.elasticsearch.action.ActionListener;
|
|
|
import org.elasticsearch.action.support.ActionFilters;
|
|
|
+import org.elasticsearch.action.support.SubscribableListener;
|
|
|
import org.elasticsearch.action.support.master.AcknowledgedResponse;
|
|
|
import org.elasticsearch.action.support.master.AcknowledgedTransportMasterNodeAction;
|
|
|
import org.elasticsearch.cluster.ClusterState;
|
|
@@ -18,6 +22,8 @@ import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
|
|
|
import org.elasticsearch.cluster.service.ClusterService;
|
|
|
import org.elasticsearch.common.inject.Inject;
|
|
|
import org.elasticsearch.common.util.concurrent.EsExecutors;
|
|
|
+import org.elasticsearch.inference.InferenceServiceRegistry;
|
|
|
+import org.elasticsearch.rest.RestStatus;
|
|
|
import org.elasticsearch.tasks.Task;
|
|
|
import org.elasticsearch.threadpool.ThreadPool;
|
|
|
import org.elasticsearch.transport.TransportService;
|
|
@@ -26,7 +32,10 @@ import org.elasticsearch.xpack.inference.registry.ModelRegistry;
|
|
|
|
|
|
public class TransportDeleteInferenceModelAction extends AcknowledgedTransportMasterNodeAction<DeleteInferenceModelAction.Request> {
|
|
|
|
|
|
+ private static final Logger logger = LogManager.getLogger(TransportPutInferenceModelAction.class);
|
|
|
+
|
|
|
private final ModelRegistry modelRegistry;
|
|
|
+ private final InferenceServiceRegistry serviceRegistry;
|
|
|
|
|
|
@Inject
|
|
|
public TransportDeleteInferenceModelAction(
|
|
@@ -35,7 +44,8 @@ public class TransportDeleteInferenceModelAction extends AcknowledgedTransportMa
|
|
|
ThreadPool threadPool,
|
|
|
ActionFilters actionFilters,
|
|
|
IndexNameExpressionResolver indexNameExpressionResolver,
|
|
|
- ModelRegistry modelRegistry
|
|
|
+ ModelRegistry modelRegistry,
|
|
|
+ InferenceServiceRegistry serviceRegistry
|
|
|
) {
|
|
|
super(
|
|
|
DeleteInferenceModelAction.NAME,
|
|
@@ -48,6 +58,7 @@ public class TransportDeleteInferenceModelAction extends AcknowledgedTransportMa
|
|
|
EsExecutors.DIRECT_EXECUTOR_SERVICE
|
|
|
);
|
|
|
this.modelRegistry = modelRegistry;
|
|
|
+ this.serviceRegistry = serviceRegistry;
|
|
|
}
|
|
|
|
|
|
@Override
|
|
@@ -57,11 +68,29 @@ public class TransportDeleteInferenceModelAction extends AcknowledgedTransportMa
|
|
|
ClusterState state,
|
|
|
ActionListener<AcknowledgedResponse> listener
|
|
|
) {
|
|
|
- modelRegistry.deleteModel(request.getModelId(), listener.delegateFailureAndWrap((l, r) -> l.onResponse(AcknowledgedResponse.TRUE)));
|
|
|
+ SubscribableListener.<ModelRegistry.UnparsedModel>newForked(modelConfigListener -> {
|
|
|
+ modelRegistry.getModel(request.getModelId(), modelConfigListener);
|
|
|
+ }).<Boolean>andThen((l1, unparsedModel) -> {
|
|
|
+ var service = serviceRegistry.getService(unparsedModel.service());
|
|
|
+ if (service.isPresent()) {
|
|
|
+ service.get().stop(request.getModelId(), l1);
|
|
|
+ } else {
|
|
|
+ l1.onFailure(new ElasticsearchStatusException("No service found for model " + request.getModelId(), RestStatus.NOT_FOUND));
|
|
|
+ }
|
|
|
+ }).<Boolean>andThen((l2, didStop) -> {
|
|
|
+ if (didStop) {
|
|
|
+ modelRegistry.deleteModel(request.getModelId(), l2);
|
|
|
+ } else {
|
|
|
+ l2.onFailure(
|
|
|
+ new ElasticsearchStatusException("Failed to stop model " + request.getModelId(), RestStatus.INTERNAL_SERVER_ERROR)
|
|
|
+ );
|
|
|
+ }
|
|
|
+ }).addListener(listener.delegateFailure((l3, didDeleteModel) -> listener.onResponse(AcknowledgedResponse.of(didDeleteModel))));
|
|
|
}
|
|
|
|
|
|
@Override
|
|
|
protected ClusterBlockException checkBlock(DeleteInferenceModelAction.Request request, ClusterState state) {
|
|
|
return state.blocks().globalBlockedException(ClusterBlockLevel.WRITE);
|
|
|
}
|
|
|
+
|
|
|
}
|