From ef8da7b43fc604ce4d8a94d8c19fd82381f96488 Mon Sep 17 00:00:00 2001
From: Yaliang Wu <ylwu@amazon.com>
Date: Wed, 11 Oct 2023 00:32:13 -0700
Subject: [PATCH] only add model info to cache if model cache exist

Signed-off-by: Yaliang Wu <ylwu@amazon.com>
---
 .../java/org/opensearch/ml/model/MLModelCacheHelper.java    | 6 ++++--
 1 file changed, 4 insertions(+), 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
index 554065ed95..553ffeb664 100644
--- 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) {