Search in sources :

Example 1 with MLModelGetRequest

use of org.opensearch.ml.common.transport.model.MLModelGetRequest in project ml-commons by opensearch-project.

the class MachineLearningNodeClient method getModel.

@Override
public void getModel(String modelId, ActionListener<MLModel> listener) {
    MLModelGetRequest mlModelGetRequest = MLModelGetRequest.builder().modelId(modelId).build();
    client.execute(MLModelGetAction.INSTANCE, mlModelGetRequest, ActionListener.wrap(response -> {
        listener.onResponse(MLModelGetResponse.fromActionResponse(response).getMlModel());
    }, listener::onFailure));
}
Also used : MLOutput(org.opensearch.ml.common.parameter.MLOutput) FieldDefaults(lombok.experimental.FieldDefaults) MLModelDeleteRequest(org.opensearch.ml.common.transport.model.MLModelDeleteRequest) RequiredArgsConstructor(lombok.RequiredArgsConstructor) MLInput(org.opensearch.ml.common.parameter.MLInput) MLTrainAndPredictionTaskAction(org.opensearch.ml.common.transport.trainpredict.MLTrainAndPredictionTaskAction) Input(org.opensearch.ml.common.parameter.Input) Function(java.util.function.Function) MLPredictionTaskRequest(org.opensearch.ml.common.transport.prediction.MLPredictionTaskRequest) Output(org.opensearch.ml.common.parameter.Output) AccessLevel(lombok.AccessLevel) MLPredictionTaskAction(org.opensearch.ml.common.transport.prediction.MLPredictionTaskAction) DeleteResponse(org.opensearch.action.delete.DeleteResponse) SearchRequest(org.opensearch.action.search.SearchRequest) ActionListener(org.opensearch.action.ActionListener) ActionResponse(org.opensearch.action.ActionResponse) SearchResponse(org.opensearch.action.search.SearchResponse) MLModel(org.opensearch.ml.common.parameter.MLModel) MLModelGetResponse(org.opensearch.ml.common.transport.model.MLModelGetResponse) MLTaskResponse(org.opensearch.ml.common.transport.MLTaskResponse) MLExecuteTaskAction(org.opensearch.ml.common.transport.execute.MLExecuteTaskAction) MLModelDeleteAction(org.opensearch.ml.common.transport.model.MLModelDeleteAction) NodeClient(org.opensearch.client.node.NodeClient) MLModelGetRequest(org.opensearch.ml.common.transport.model.MLModelGetRequest) MLModelSearchAction(org.opensearch.ml.common.transport.model.MLModelSearchAction) MLExecuteTaskRequest(org.opensearch.ml.common.transport.execute.MLExecuteTaskRequest) MLTrainingTaskAction(org.opensearch.ml.common.transport.training.MLTrainingTaskAction) MLModelGetAction(org.opensearch.ml.common.transport.model.MLModelGetAction) MLExecuteTaskResponse(org.opensearch.ml.common.transport.execute.MLExecuteTaskResponse) MLTrainingTaskRequest(org.opensearch.ml.common.transport.training.MLTrainingTaskRequest) MLModelGetRequest(org.opensearch.ml.common.transport.model.MLModelGetRequest)

Example 2 with MLModelGetRequest

use of org.opensearch.ml.common.transport.model.MLModelGetRequest in project ml-commons by opensearch-project.

the class GetModelTransportAction method doExecute.

@Override
protected void doExecute(Task task, ActionRequest request, ActionListener<MLModelGetResponse> actionListener) {
    MLModelGetRequest mlModelGetRequest = MLModelGetRequest.fromActionRequest(request);
    String modelId = mlModelGetRequest.getModelId();
    GetRequest getRequest = new GetRequest(ML_MODEL_INDEX).id(modelId);
    try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
        client.get(getRequest, ActionListener.wrap(r -> {
            log.info("Completed Get Model Request, id:{}", modelId);
            if (r != null && r.isExists()) {
                try (XContentParser parser = createXContentParserFromRegistry(xContentRegistry, r.getSourceAsBytesRef())) {
                    ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
                    MLModel mlModel = MLModel.parse(parser);
                    actionListener.onResponse(MLModelGetResponse.builder().mlModel(mlModel).build());
                } catch (Exception e) {
                    log.error("Failed to parse ml model" + r.getId(), e);
                    actionListener.onFailure(e);
                }
            } else {
                actionListener.onFailure(new MLResourceNotFoundException("Fail to find model"));
            }
        }, e -> {
            if (e instanceof IndexNotFoundException) {
                actionListener.onFailure(new MLResourceNotFoundException("Fail to find model"));
            } else {
                log.error("Failed to get ML model " + modelId, e);
                actionListener.onFailure(e);
            }
        }));
    } catch (Exception e) {
        log.error("Failed to get ML model " + modelId, e);
        actionListener.onFailure(e);
    }
}
Also used : FieldDefaults(lombok.experimental.FieldDefaults) Client(org.opensearch.client.Client) HandledTransportAction(org.opensearch.action.support.HandledTransportAction) MLModelGetRequest(org.opensearch.ml.common.transport.model.MLModelGetRequest) IndexNotFoundException(org.opensearch.index.IndexNotFoundException) GetRequest(org.opensearch.action.get.GetRequest) XContentParserUtils.ensureExpectedToken(org.opensearch.common.xcontent.XContentParserUtils.ensureExpectedToken) MLNodeUtils.createXContentParserFromRegistry(org.opensearch.ml.utils.MLNodeUtils.createXContentParserFromRegistry) ActionRequest(org.opensearch.action.ActionRequest) Task(org.opensearch.tasks.Task) ThreadContext(org.opensearch.common.util.concurrent.ThreadContext) TransportService(org.opensearch.transport.TransportService) MLResourceNotFoundException(org.opensearch.ml.common.exception.MLResourceNotFoundException) XContentParser(org.opensearch.common.xcontent.XContentParser) ActionFilters(org.opensearch.action.support.ActionFilters) AccessLevel(lombok.AccessLevel) MLModelGetAction(org.opensearch.ml.common.transport.model.MLModelGetAction) ML_MODEL_INDEX(org.opensearch.ml.indices.MLIndicesHandler.ML_MODEL_INDEX) NamedXContentRegistry(org.opensearch.common.xcontent.NamedXContentRegistry) Log4j2(lombok.extern.log4j.Log4j2) Inject(org.opensearch.common.inject.Inject) ActionListener(org.opensearch.action.ActionListener) MLModel(org.opensearch.ml.common.parameter.MLModel) MLModelGetResponse(org.opensearch.ml.common.transport.model.MLModelGetResponse) MLResourceNotFoundException(org.opensearch.ml.common.exception.MLResourceNotFoundException) MLModelGetRequest(org.opensearch.ml.common.transport.model.MLModelGetRequest) GetRequest(org.opensearch.action.get.GetRequest) ThreadContext(org.opensearch.common.util.concurrent.ThreadContext) IndexNotFoundException(org.opensearch.index.IndexNotFoundException) MLModel(org.opensearch.ml.common.parameter.MLModel) MLModelGetRequest(org.opensearch.ml.common.transport.model.MLModelGetRequest) XContentParser(org.opensearch.common.xcontent.XContentParser) IndexNotFoundException(org.opensearch.index.IndexNotFoundException) MLResourceNotFoundException(org.opensearch.ml.common.exception.MLResourceNotFoundException)

Aggregations

AccessLevel (lombok.AccessLevel)2 FieldDefaults (lombok.experimental.FieldDefaults)2 ActionListener (org.opensearch.action.ActionListener)2 MLModel (org.opensearch.ml.common.parameter.MLModel)2 MLModelGetAction (org.opensearch.ml.common.transport.model.MLModelGetAction)2 MLModelGetRequest (org.opensearch.ml.common.transport.model.MLModelGetRequest)2 MLModelGetResponse (org.opensearch.ml.common.transport.model.MLModelGetResponse)2 Function (java.util.function.Function)1 RequiredArgsConstructor (lombok.RequiredArgsConstructor)1 Log4j2 (lombok.extern.log4j.Log4j2)1 ActionRequest (org.opensearch.action.ActionRequest)1 ActionResponse (org.opensearch.action.ActionResponse)1 DeleteResponse (org.opensearch.action.delete.DeleteResponse)1 GetRequest (org.opensearch.action.get.GetRequest)1 SearchRequest (org.opensearch.action.search.SearchRequest)1 SearchResponse (org.opensearch.action.search.SearchResponse)1 ActionFilters (org.opensearch.action.support.ActionFilters)1 HandledTransportAction (org.opensearch.action.support.HandledTransportAction)1 Client (org.opensearch.client.Client)1 NodeClient (org.opensearch.client.node.NodeClient)1