Examples with MLPredictionOutput org.opensearch.ml.common.parameter.MLPredictionOutput used on opensource projects

Example 1 with MLPredictionOutput

use of org.opensearch.ml.common.parameter.MLPredictionOutput in project ml-commons by opensearch-project.

the class MLTrainAndPredictTaskRunner method trainAndPredict.

private void trainAndPredict(MLTask mlTask, DataFrame inputDataFrame, MLTrainingTaskRequest request, ActionListener<MLTaskResponse> listener) {
    ActionListener<MLTaskResponse> internalListener = wrappedCleanupListener(listener, mlTask.getTaskId());
    // track ML task count and add ML task into cache
    mlStats.getStat(ML_EXECUTING_TASK_COUNT).increment();
    mlStats.getStat(ML_TOTAL_REQUEST_COUNT).increment();
    mlStats.createCounterStatIfAbsent(requestCountStat(mlTask.getFunctionName(), ActionName.TRAIN_PREDICT)).increment();
    mlTaskManager.add(mlTask);
    MLInput mlInput = request.getMlInput();
    // run train and predict
    try {
        mlTaskManager.updateTaskState(mlTask.getTaskId(), MLTaskState.RUNNING, mlTask.isAsync());
        MLOutput output = MLEngine.trainAndPredict(mlInput.toBuilder().inputDataset(new DataFrameInputDataset(inputDataFrame)).build());
        handleAsyncMLTaskComplete(mlTask);
        if (output instanceof MLPredictionOutput) {
            ((MLPredictionOutput) output).setStatus(MLTaskState.COMPLETED.name());
        }
        MLTaskResponse response = MLTaskResponse.builder().output(output).build();
        log.info("Train and predict task done for algorithm: {}, task id: {}", mlTask.getFunctionName(), mlTask.getTaskId());
        internalListener.onResponse(response);
    } catch (Exception e) {
        // todo need to specify what exception
        log.error("Failed to train and predict " + mlInput.getAlgorithm(), e);
        handlePredictFailure(mlTask, listener, e, true);
        return;
    }
}

Example 2 with MLPredictionOutput

use of org.opensearch.ml.common.parameter.MLPredictionOutput in project ml-commons by opensearch-project.

the class MLPredictTaskRunner method predict.

private void predict(MLTask mlTask, DataFrame inputDataFrame, MLPredictionTaskRequest request, ActionListener<MLTaskResponse> listener) {
    ActionListener<MLTaskResponse> internalListener = wrappedCleanupListener(listener, mlTask.getTaskId());
    // track ML task count and add ML task into cache
    mlStats.getStat(ML_EXECUTING_TASK_COUNT).increment();
    mlStats.getStat(ML_TOTAL_REQUEST_COUNT).increment();
    mlStats.createCounterStatIfAbsent(requestCountStat(mlTask.getFunctionName(), ActionName.PREDICT)).increment();
    mlTaskManager.add(mlTask);
    // run predict
    if (request.getModelId() != null) {
        // search model by model id.
        try (ThreadContext.StoredContext context = threadPool.getThreadContext().stashContext()) {
            MLInput mlInput = request.getMlInput();
            ActionListener<GetResponse> getResponseListener = ActionListener.wrap(r -> {
                if (r == null || !r.isExists()) {
                    internalListener.onFailure(new ResourceNotFoundException("No model found, please check the modelId."));
                    return;
                }
                Map<String, Object> source = r.getSourceAsMap();
                User requestUser = getUserContext(client);
                User resourceUser = User.parse((String) source.get(USER));
                if (!checkUserPermissions(requestUser, resourceUser, request.getModelId())) {
                    // The backend roles of request user and resource user doesn't have intersection
                    OpenSearchException e = new OpenSearchException("User: " + requestUser.getName() + " does not have permissions to run predict by model: " + request.getModelId());
                    handlePredictFailure(mlTask, internalListener, e, false);
                    return;
                }
                Model model = new Model();
                model.setName((String) source.get(MLModel.MODEL_NAME));
                model.setVersion((Integer) source.get(MLModel.MODEL_VERSION));
                byte[] decoded = Base64.getDecoder().decode((String) source.get(MLModel.MODEL_CONTENT));
                model.setContent(decoded);
                // run predict
                mlTaskManager.updateTaskState(mlTask.getTaskId(), MLTaskState.RUNNING, mlTask.isAsync());
                MLOutput output = MLEngine.predict(mlInput.toBuilder().inputDataset(new DataFrameInputDataset(inputDataFrame)).build(), model);
                if (output instanceof MLPredictionOutput) {
                    ((MLPredictionOutput) output).setStatus(MLTaskState.COMPLETED.name());
                }
                // Once prediction complete, reduce ML_EXECUTING_TASK_COUNT and update task state
                handleAsyncMLTaskComplete(mlTask);
                MLTaskResponse response = MLTaskResponse.builder().output(output).build();
                internalListener.onResponse(response);
            }, e -> {
                log.error("Failed to predict " + mlInput.getAlgorithm() + ", modelId: " + mlTask.getModelId(), e);
                handlePredictFailure(mlTask, internalListener, e, true);
            });
            GetRequest getRequest = new GetRequest(ML_MODEL_INDEX, mlTask.getModelId());
            client.get(getRequest, ActionListener.runBefore(getResponseListener, () -> context.restore()));
        } catch (Exception e) {
            log.error("Failed to get model " + mlTask.getModelId(), e);
            handlePredictFailure(mlTask, internalListener, e, true);
        }
    } else {
        IllegalArgumentException e = new IllegalArgumentException("ModelId is invalid");
        log.error("ModelId is invalid", e);
        handlePredictFailure(mlTask, internalListener, e, false);
    }
}

Example 3 with MLPredictionOutput

use of org.opensearch.ml.common.parameter.MLPredictionOutput in project ml-commons by opensearch-project.

the class KMeansTest method predict.

@Test
public void predict() {
    Model model = kMeans.train(trainDataFrame);
    MLPredictionOutput output = (MLPredictionOutput) kMeans.predict(predictionDataFrame, model);
    DataFrame predictions = output.getPredictionResult();
    Assert.assertEquals(predictionSize, predictions.size());
    predictions.forEach(row -> Assert.assertTrue(row.getValue(0).intValue() == 0 || row.getValue(0).intValue() == 1));
}

Example 4 with MLPredictionOutput

use of org.opensearch.ml.common.parameter.MLPredictionOutput in project ml-commons by opensearch-project.

the class BatchRandomCutForestTest method predictWithNullModel.

@Test
public void predictWithNullModel() {
    exceptionRule.expect(IllegalArgumentException.class);
    exceptionRule.expectMessage("No model found for batch RCF prediction");
    MLPredictionOutput output = (MLPredictionOutput) forest.predict(predictionDataFrame, null);
    verifyPredictionResult(output);
}

Example 5 with MLPredictionOutput

use of org.opensearch.ml.common.parameter.MLPredictionOutput in project ml-commons by opensearch-project.

the class MLPredictionTaskResponseTest method toXContentTest.

@Test
public void toXContentTest() throws IOException {
    MLPredictionOutput output = MLPredictionOutput.builder().taskId("b5009b99-268f-476d-a676-379a30f82457").status("Success").predictionResult(DataFrameBuilder.load(Collections.singletonList(new HashMap<String, Object>() {

        {
            put("ClusterID", 0);
        }
    }))).build();
    MLTaskResponse response = MLTaskResponse.builder().output(output).build();
    XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
    response.toXContent(builder, ToXContent.EMPTY_PARAMS);
    assertNotNull(builder);
    String jsonStr = Strings.toString(builder);
    assertEquals("{\"task_id\":\"b5009b99-268f-476d-a676-379a30f82457\"," + "\"status\":\"Success\"," + "\"prediction_result\":{" + "\"column_metas\":[{\"name\":\"ClusterID\",\"column_type\":\"INTEGER\"}]," + "\"rows\":[{\"values\":[{\"column_type\":\"INTEGER\",\"value\":0}]}]}}", jsonStr);
}