Fix model not deploy issue under intensive prediction taskks

Signed-off-by: zane-neo <[email protected]>
zane-neo · Jan 23, 2024 · 9d7b9dc · 9d7b9dc
1 parent de59efc
commit 9d7b9dc
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/plugin/src/main/java/org/opensearch/ml/action/syncup/TransportSyncUpOnNodeAction.java b/plugin/src/main/java/org/opensearch/ml/action/syncup/TransportSyncUpOnNodeAction.java
@@ -198,6 +198,10 @@ void cleanUpLocalCache(Map<String, Set<String>> runningDeployModelTasks) {
         }
         for (String taskId : allTaskIds) {
             MLTaskCache mlTaskCache = mlTaskManager.getMLTaskCache(taskId);
+            // Task could be a prediction task, and it could be completed and removed from cache in predict thread during the cleaning up.
+            if (mlTaskCache == null) {
+                continue;
+            }
             MLTask mlTask = mlTaskCache.getMlTask();
             Instant lastUpdateTime = mlTask.getLastUpdateTime();
             Instant now = Instant.now();

diff --git a/plugin/src/main/java/org/opensearch/ml/task/MLTaskManager.java b/plugin/src/main/java/org/opensearch/ml/task/MLTaskManager.java
@@ -117,7 +117,7 @@ public synchronized void add(MLTask mlTask, List<String> workerNodes) {
             throw new IllegalArgumentException("Duplicate taskId");
         }
         taskCaches.put(taskId, new MLTaskCache(mlTask, workerNodes));
-        log.debug("add ML task to cache " + taskId);
+        log.debug("add ML task to cache, taskId: {}, taskType: {} ", taskId, mlTask.getTaskType());
     }
 
     /**