fix no worker node exception for remote embedding model (#1482) (#1511)

* fix no worker node exception for remote embedding model Signed-off-by: Yaliang Wu <ylwu@amazon.com> * only add model info to cache if model cache exist Signed-off-by: Yaliang Wu <ylwu@amazon.com> --------- Signed-off-by: Yaliang Wu <ylwu@amazon.com> (cherry picked from commit 6f83b9f) Co-authored-by: Yaliang Wu <ylwu@amazon.com>
opensearch-project · Oct 12, 2023 · 0a08655 · 0a08655
1 parent e4851d2
commit 0a08655
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java b/plugin/src/main/java/org/opensearch/ml/action/prediction/TransportPredictionTaskAction.java
@@ -137,9 +137,16 @@ private void executePredict(
         String requestId = mlPredictionTaskRequest.getRequestID();
         log.debug("receive predict request " + requestId + " for model " + mlPredictionTaskRequest.getModelId());
         long startTime = System.nanoTime();
+        // For remote text embedding model, neural search will set mlPredictionTaskRequest.getMlInput().getAlgorithm() as
+        // TEXT_EMBEDDING. In ml-commons we should always use the real function name of model: REMOTE. So we try to get
+        // from model cache first.
+        FunctionName functionName = modelCacheHelper
+            .getOptionalFunctionName(modelId)
+            .orElse(mlPredictionTaskRequest.getMlInput().getAlgorithm());
         mlPredictTaskRunner
             .run(
-                mlPredictionTaskRequest.getMlInput().getAlgorithm(),
+                // This is by design to NOT use mlPredictionTaskRequest.getMlInput().getAlgorithm() here
+                functionName,
                 mlPredictionTaskRequest,
                 transportService,
                 ActionListener.runAfter(wrappedListener, () -> {

diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelCacheHelper.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelCacheHelper.java
@@ -431,8 +431,10 @@ public boolean getDeployToAllNodes(String modelId) {
     }
 
     public void setModelInfo(String modelId, MLModel mlModel) {
-        MLModelCache mlModelCache = getExistingModelCache(modelId);
-        mlModelCache.setModelInfo(mlModel);
+        MLModelCache mlModelCache = modelCaches.get(modelId);
+        if (mlModelCache != null) {
+            mlModelCache.setModelInfo(mlModel);
+        }
     }
 
     public MLModel getModelInfo(String modelId) {