only add model info to cache if model cache exist

Signed-off-by: Yaliang Wu <[email protected]>
opensearch-project · Oct 11, 2023 · ef8da7b · ef8da7b
1 parent 99e34a0
commit ef8da7b
Showing 1 changed file with 4 additions and 2 deletions.
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) {