debugging

Signed-off-by: Sicheng Song <sicheng.song@outlook.com>

Author: b4sjoo

plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java

@@ -236,12 +236,12 @@ private void executePredict(
             );
     }
 
-    private void validateInputSchema(String modelId, MLInput mlInput) throws IOException {
+    private void validateInputSchema(String modelId, MLInput mlInput) {
         if (modelCacheHelper.getModelInterface(modelId) != null && modelCacheHelper.getModelInterface(modelId).get("input") != null) {
             String inputSchemaString = modelCacheHelper.getModelInterface(modelId).get("input");
-            log.info(inputSchemaString);
-            log.info(mlInput.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS).toString());
             try {
+                log.info(inputSchemaString);
+                log.info(mlInput.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS).toString());
                 MLNodeUtils
                     .validateSchema(
                         inputSchemaString,

plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java

@@ -14,6 +14,7 @@
 import static org.opensearch.ml.plugin.MachineLearningPlugin.REMOTE_PREDICT_THREAD_POOL;
 import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_MODEL_AUTO_DEPLOY_ENABLE;
 
+import java.io.IOException;
 import java.time.Instant;
 import java.util.Arrays;
 import java.util.UUID;
@@ -330,7 +331,20 @@ private void runPredict(
                     if (output instanceof MLPredictionOutput) {
                         ((MLPredictionOutput) output).setStatus(MLTaskState.COMPLETED.name());
                     }
-
+                    if (output instanceof ModelTensorOutput) {
+                        log.info("ModelId " + modelId + " is completed as ModelTensorOutput.");
+                        // pick the first output tensor to validate the schema
+                        if (((ModelTensorOutput) output).getMlModelOutputs() != null
+                                && !((ModelTensorOutput) output).getMlModelOutputs().isEmpty()) {
+                            ModelTensors modelTensors = ((ModelTensorOutput) output).getMlModelOutputs().get(0);
+                            log.info(modelTensors.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS).toString());
+                            if (modelTensors.getMlModelTensors() != null && !modelTensors.getMlModelTensors().isEmpty()) {
+                                ModelTensor modelTensorOutput = modelTensors.getMlModelTensors().get(0);
+                                log.info(modelTensorOutput.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS).toString());
+                                validateOutputSchema(modelId, modelTensorOutput);
+                            }
+                        }
+                    }
                     // Once prediction complete, reduce ML_EXECUTING_TASK_COUNT and update task state
                     handleAsyncMLTaskComplete(mlTask);
                     MLTaskResponse response = MLTaskResponse.builder().output(output).build();
@@ -377,7 +391,7 @@ private void runPredict(
                         }
                         MLOutput output = mlEngine.predict(mlInput, mlModel);
                         if (output instanceof MLPredictionOutput) {
-                            log.info("ModelId " + modelId + " is completed as MLPredictionOutput.");
+                            ((MLPredictionOutput) output).setStatus(MLTaskState.COMPLETED.name());
                         }
                         if (output instanceof ModelTensorOutput) {
                             log.info("ModelId " + modelId + " is completed as ModelTensorOutput.");
@@ -453,8 +467,9 @@ private void validateOutputSchema(String modelId, ModelTensor modelTensor) {
         log.info("Went to validate output schema");
         if (mlModelManager.getModelInterface(modelId) != null && mlModelManager.getModelInterface(modelId).get("output") != null) {
             String outputSchemaString = mlModelManager.getModelInterface(modelId).get("output");
-            log.info(outputSchemaString);
             try {
+                log.info(outputSchemaString);
+                log.info(modelTensor.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS).toString());
                 MLNodeUtils
                     .validateSchema(
                         outputSchemaString,