check model auto deploy (#2288) (#2290)

Signed-off-by: Xun Zhang <[email protected]> (cherry picked from commit 848abf6) Co-authored-by: Xun Zhang <[email protected]>
opensearch-project · Mar 27, 2024 · 493e394 · 493e394
1 parent c5c4e61
commit 493e394
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
@@ -148,6 +148,16 @@ public void dispatchTask(
                 if (FunctionName.isAutoDeployEnabled(autoDeploymentEnabled, functionName)) {
                     try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
                         mlModelManager.getModel(modelId, ActionListener.runBefore(ActionListener.wrap(model -> {
+                            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;
+                            }
                             String[] planningWorkerNodes = model.getPlanningWorkerNodes();
                             MLModel modelBeingAutoDeployed = mlModelManager.addModelToAutoDeployCache(modelId, model);
                             if (modelBeingAutoDeployed == model) {
@@ -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;
     }