debugging

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

Author: b4sjoo

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

@@ -45,6 +45,8 @@
 import lombok.experimental.FieldDefaults;
 import lombok.extern.log4j.Log4j2;
 
+import java.io.IOException;
+
 @Log4j2
 @FieldDefaults(level = AccessLevel.PRIVATE)
 public class TransportPredictionTaskAction extends HandledTransportAction<ActionRequest, MLTaskResponse> {
@@ -160,8 +162,6 @@ public void onResponse(MLModel mlModel) {
                                             executePredict(mlPredictionTaskRequest, wrappedListener, modelId);
                                         }
                                     } else {
-                                        log.info("ModelId " + modelId + " is enabled for prediction in remote layer.");
-                                        log.info(mlPredictionTaskRequest.getMlInput().toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS).toString());
                                         validateInputSchema(modelId, mlPredictionTaskRequest.getMlInput());
                                         executePredict(mlPredictionTaskRequest, wrappedListener, modelId);
                                     }
@@ -236,21 +236,19 @@ private void executePredict(
             );
     }
 
-    private void validateInputSchema(String modelId, MLInput mlInput) {
-        log.info("Went to validate input schema");
-        log.info(modelCacheHelper.getModelInterface(modelId).containsKey("input"));
-        log.info(modelCacheHelper.getModelInterface(modelId).get("input"));
+    private void validateInputSchema(String modelId, MLInput mlInput) throws IOException {
         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 {
                 MLNodeUtils
                     .validateSchema(
                         inputSchemaString,
                         mlInput.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS).toString()
                     );
             } catch (Exception e) {
-                throw new IllegalArgumentException("Error validating input schema: " + e.getMessage());
+                throw new OpenSearchStatusException("Error validating input schema: " + e.getMessage(), RestStatus.BAD_REQUEST);
             }
         }
     }

plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java

@@ -1486,6 +1486,7 @@ public Map<String, TokenBucket> getUserRateLimiterMap(String modelId) {
      * Set up model interface with model id.
      */
     private void setupModelInterface(String modelId, Map<String, String> modelInterface) {
+        log.debug("Model interface for model: {} loaded into cache.", modelId);
         if (modelInterface != null) {
             modelCacheHelper.setModelInterface(modelId, modelInterface);
         } else {

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

@@ -377,10 +377,10 @@ private void runPredict(
                         }
                         MLOutput output = mlEngine.predict(mlInput, mlModel);
                         if (output instanceof MLPredictionOutput) {
-                            ((MLPredictionOutput) output).setStatus(MLTaskState.COMPLETED.name());
+                            log.info("ModelId " + modelId + " is completed as MLPredictionOutput.");
                         }
                         if (output instanceof ModelTensorOutput) {
-                            log.info("ModelId " + modelId + " is completed for prediction.");
+                            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()) {