set the mlModel info in the deploy stage (#2389) (#2391)

Signed-off-by: Xun Zhang <[email protected]> (cherry picked from commit ff6048f) Co-authored-by: Xun Zhang <[email protected]>
opensearch-project · May 1, 2024 · c873f22 · c873f22
1 parent 0e9a4b2
commit c873f22
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 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
@@ -386,9 +386,11 @@ public String[] getLocalDeployedModels() {
      */
     public String[] getExpiredModels() {
         return modelCaches.entrySet().stream().filter(entry -> {
-            MLModel mlModel = entry.getValue().getCachedModelInfo();
-            if (mlModel.getDeploySetting() == null) {
-                return false;  // no TTL, never expire
+            MLModelCache modelCache = entry.getValue();
+            MLModel mlModel = modelCache.getCachedModelInfo();
+            MLModelState modelState = modelCache.getModelState();
+            if (mlModel == null || mlModel.getDeploySetting() == null) {
+                return false; // no TTL, never expire
             }
             Duration liveDuration = Duration.between(entry.getValue().getLastAccessTime(), Instant.now());
             Long ttlInMinutes = mlModel.getDeploySetting().getModelTTLInMinutes();
@@ -397,8 +399,7 @@ public String[] getExpiredModels() {
             }
             Duration ttl = Duration.ofMinutes(ttlInMinutes);
             boolean isModelExpired = liveDuration.getSeconds() >= ttl.getSeconds();
-            return isModelExpired
-                && (mlModel.getModelState() == MLModelState.DEPLOYED || mlModel.getModelState() == MLModelState.PARTIALLY_DEPLOYED);
+            return isModelExpired && (modelState == MLModelState.DEPLOYED || modelState == MLModelState.PARTIALLY_DEPLOYED);
         }).map(entry -> entry.getKey()).collect(Collectors.toList()).toArray(new String[0]);
     }
 

diff --git a/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java b/plugin/src/main/java/org/opensearch/ml/model/MLModelManager.java
@@ -996,6 +996,7 @@ public void deployModel(
             }
             this.getModel(modelId, threadedActionListener(DEPLOY_THREAD_POOL, ActionListener.wrap(mlModel -> {
                 modelCacheHelper.setIsModelEnabled(modelId, mlModel.getIsEnabled());
+                modelCacheHelper.setModelInfo(modelId, mlModel);
                 if (FunctionName.REMOTE == mlModel.getAlgorithm()
                     || (!FunctionName.isDLModel(mlModel.getAlgorithm()) && mlModel.getAlgorithm() != FunctionName.METRICS_CORRELATION)) {
                     // deploy remote model or model trained by built-in algorithm like kmeans