From 103633aa7b99d8d76b6b50ec6bb18428ec1160ad Mon Sep 17 00:00:00 2001
From: David Kyle <david.kyle@elastic.co>
Date: Fri, 24 May 2024 11:40:30 +0100
Subject: [PATCH] [ML]  Fix NPE in trained model assignment updater (#108942)

The NPE causes the rebalance to fail, later retires would have corrected
the situation
---
 docs/changelog/108942.yaml                    |  5 ++++
 .../TrainedModelAssignmentClusterService.java | 23 ++++++++++++-------
 2 files changed, 20 insertions(+), 8 deletions(-)
 create mode 100644 docs/changelog/108942.yaml

diff --git a/docs/changelog/108942.yaml b/docs/changelog/108942.yaml
new file mode 100644
index 000000000000..c58b06a92cee
--- /dev/null
+++ b/docs/changelog/108942.yaml
@@ -0,0 +1,5 @@
+pr: 108942
+summary: Fix NPE in trained model assignment updater
+area: Machine Learning
+type: bug
+issues: []
diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentClusterService.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentClusterService.java
index 637ad9d7bbbb..f468e5239fd2 100644
--- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentClusterService.java
+++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentClusterService.java
@@ -1047,15 +1047,22 @@ public class TrainedModelAssignmentClusterService implements ClusterStateListene
         if (event.changedCustomMetadataSet().contains(PersistentTasksCustomMetadata.TYPE) == false) {
             return Optional.empty();
         }
-        final PersistentTasksCustomMetadata previousPersistentTasks = event.previousState()
-            .getMetadata()
-            .custom(PersistentTasksCustomMetadata.TYPE);
-        final PersistentTasksCustomMetadata currentPersistentTasks = event.state().getMetadata().custom(PersistentTasksCustomMetadata.TYPE);
-        Set<String> previousMlTaskIds = findMlProcessTaskIds(previousPersistentTasks);
+
+        PersistentTasksCustomMetadata previousPersistentTasks = PersistentTasksCustomMetadata.getPersistentTasksCustomMetadata(
+            event.previousState()
+        );
+        if (previousPersistentTasks == null) { // no previous jobs so nothing has stopped
+            return Optional.empty();
+        }
+
+        PersistentTasksCustomMetadata currentPersistentTasks = PersistentTasksCustomMetadata.getPersistentTasksCustomMetadata(
+            event.state()
+        );
         Set<String> currentMlTaskIds = findMlProcessTaskIds(currentPersistentTasks);
-        Set<String> stoppedTaskTypes = previousMlTaskIds.stream()
-            .filter(id -> currentMlTaskIds.contains(id) == false) // remove the tasks that are still present. Stopped Ids only.
-            .map(previousPersistentTasks::getTask)
+
+        Set<PersistentTasksCustomMetadata.PersistentTask<?>> previousMlTasks = MlTasks.findMlProcessTasks(previousPersistentTasks);
+        Set<String> stoppedTaskTypes = previousMlTasks.stream()
+            .filter(task -> currentMlTaskIds.contains(task.getId()) == false) // remove the tasks that are still present. Stopped Ids only.
             .map(PersistentTasksCustomMetadata.PersistentTask::getTaskName)
             .map(MlTasks::prettyPrintTaskName)
             .collect(Collectors.toSet());