From c873f221e0dcd0cd6ca6189b7cdabc9ff1368f74 Mon Sep 17 00:00:00 2001 From: "opensearch-trigger-bot[bot]" <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Date: Tue, 30 Apr 2024 21:52:27 -0700 Subject: [PATCH] set the mlModel info in the deploy stage (#2389) (#2391) Signed-off-by: Xun Zhang (cherry picked from commit ff6048f47c79042fcebc03cdb4196b0474d694e7) Co-authored-by: Xun Zhang --- .../org/opensearch/ml/model/MLModelCacheHelper.java | 11 ++++++----- .../java/org/opensearch/ml/model/MLModelManager.java | 1 + 2 files changed, 7 insertions(+), 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 index 6c4b12fcf4..6230ad2944 100644 --- 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 index d46dca0012..d8ba096e21 100644 --- 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