Performance enhacement for predict action by caching model info (open…

…search-project#1472) * Performance enhacement for predict action by caching model info Signed-off-by: zane-neo <[email protected]> * Add context.restore() to avoid missing info Signed-off-by: zane-neo <[email protected]> --------- Signed-off-by: zane-neo <[email protected]> (cherry picked from commit a985f6e)
rbhavna · Oct 11, 2023 · 570547b · 570547b
1 parent da5d829
commit 570547b
Showing 1 changed file with 1 addition and 4 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
@@ -436,10 +436,7 @@ public void setModelInfo(String modelId, MLModel mlModel) {
     }
 
     public MLModel getModelInfo(String modelId) {
-        MLModelCache mlModelCache = modelCaches.get(modelId);
-        if (mlModelCache == null) {
-            return null;
-        }
+        MLModelCache mlModelCache = getExistingModelCache(modelId);
         return mlModelCache.getCachedModelInfo();
     }