From c0e1a2136bb66eaef4a0680e2fb447a9e4eab25c Mon Sep 17 00:00:00 2001
From: Xun Zhang <xunzh@amazon.com>
Date: Wed, 27 Mar 2024 15:14:36 -0700
Subject: [PATCH] check model auto deploy

Signed-off-by: Xun Zhang <xunzh@amazon.com>
---
 .../opensearch/ml/task/MLPredictTaskRunner.java | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java b/plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java
index a6d82e979f..39bd8c84a0 100644
--- a/plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java
+++ b/plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java
@@ -149,6 +149,16 @@ public void dispatchTask(
                     try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
                         mlModelManager.getModel(modelId, ActionListener.runBefore(ActionListener.wrap(model -> {
                             String[] planningWorkerNodes = model.getPlanningWorkerNodes();
+                            if (!checkModelAutoDeployEnabled(model)) {
+                                log.info("Auto deployment disabled for model {}, please deploy model first", modelId);
+                                listener
+                                    .onFailure(
+                                        new IllegalArgumentException(
+                                            "Model not ready yet. Please run this first: POST /_plugins/_ml/models/" + modelId + "/_deploy"
+                                        )
+                                    );
+                                return;
+                            }
                             MLModel modelBeingAutoDeployed = mlModelManager.addModelToAutoDeployCache(modelId, model);
                             if (modelBeingAutoDeployed == model) {
                                 log.info("Automatically deploy model {}", modelId);
@@ -241,6 +251,13 @@ protected void executeTask(MLPredictionTaskRequest request, ActionListener<MLTas
         }
     }
 
+    private boolean checkModelAutoDeployEnabled(MLModel mlModel) {
+        if (mlModel.getDeploySetting() == null) {
+            return true;
+        }
+        return mlModel.getDeploySetting().getIsAutoDeployEnabled();
+    }
+
     private String getPredictThreadPool(FunctionName functionName) {
         return functionName == FunctionName.REMOTE ? REMOTE_PREDICT_THREAD_POOL : PREDICT_THREAD_POOL;
     }