fix no worker node exception for remote embedding model (opensearch-p…

…roject#1482) * fix no worker node exception for remote embedding model Signed-off-by: Yaliang Wu <[email protected]> * only add model info to cache if model cache exist Signed-off-by: Yaliang Wu <[email protected]> --------- Signed-off-by: Yaliang Wu <[email protected]>
austintlee · Oct 11, 2023 · 6f83b9f · 6f83b9f
1 parent da5d829
commit 6f83b9f
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) {