check model auto deploy

Signed-off-by: Xun Zhang <xunzh@amazon.com>
opensearch-project · Mar 27, 2024 · c0e1a21 · c0e1a21
1 parent 22fcaf0
commit c0e1a21
Showing 1 changed file with 17 additions and 0 deletions.
diff --git 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;
     }