[ML] Release cluster state (#136769) (#136833)

Refactoring TransportMlMemoryAction to not retain cluster state through
the lifecycle of the request so that it can be garbage collected as soon
as possible.

Fix #123243

Author: prwhelan

docs/changelog/136769.yaml

@@ -0,0 +1,6 @@
+pr: 136769
+summary: Release cluster state
+area: Machine Learning
+type: bug
+issues:
+ - 123243

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/assignment/TrainedModelAssignmentMetadata.java

@@ -17,6 +17,7 @@
 import org.elasticsearch.cluster.NamedDiff;
 import org.elasticsearch.cluster.SimpleDiffable;
 import org.elasticsearch.cluster.metadata.Metadata;
+import org.elasticsearch.cluster.metadata.ProjectMetadata;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -80,6 +81,14 @@ public static TrainedModelAssignmentMetadata fromState(ClusterState clusterState
         return trainedModelAssignmentMetadata == null ? EMPTY : trainedModelAssignmentMetadata;
     }
 
+    public static TrainedModelAssignmentMetadata fromMetadata(ProjectMetadata projectMetadata) {
+        TrainedModelAssignmentMetadata trainedModelAssignmentMetadata = projectMetadata.custom(NAME);
+        if (trainedModelAssignmentMetadata == null) {
+            trainedModelAssignmentMetadata = projectMetadata.custom(DEPRECATED_NAME);
+        }
+        return trainedModelAssignmentMetadata == null ? EMPTY : trainedModelAssignmentMetadata;
+    }
+
     public static List<TrainedModelAssignment> assignmentsForModelId(ClusterState clusterState, String modelId) {
         return TrainedModelAssignmentMetadata.fromState(clusterState)
             .allAssignments()

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportMlMemoryAction.java

@@ -16,11 +16,13 @@
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.client.internal.OriginSettingClient;
 import org.elasticsearch.client.internal.ParentTaskAssigningClient;
+import org.elasticsearch.cluster.ClusterName;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.block.ClusterBlockException;
 import org.elasticsearch.cluster.block.ClusterBlockLevel;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodeRole;
+import org.elasticsearch.cluster.project.ProjectResolver;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.unit.ByteSizeValue;
@@ -35,6 +37,7 @@
 import org.elasticsearch.xpack.core.ml.action.MlMemoryAction.Response.MlMemoryStats;
 import org.elasticsearch.xpack.core.ml.action.TrainedModelCacheInfoAction;
 import org.elasticsearch.xpack.core.ml.action.TrainedModelCacheInfoAction.Response.CacheInfo;
+import org.elasticsearch.xpack.core.ml.inference.assignment.TrainedModelAssignmentMetadata;
 import org.elasticsearch.xpack.ml.job.NodeLoad;
 import org.elasticsearch.xpack.ml.job.NodeLoadDetector;
 import org.elasticsearch.xpack.ml.process.MlMemoryTracker;
@@ -53,6 +56,7 @@ public class TransportMlMemoryAction extends TransportMasterNodeAction<MlMemoryA
 
     private final Client client;
     private final MlMemoryTracker memoryTracker;
+    private final ProjectResolver projectResolver;
 
     @Inject
     public TransportMlMemoryAction(
@@ -61,7 +65,8 @@ public TransportMlMemoryAction(
         ThreadPool threadPool,
         ActionFilters actionFilters,
         Client client,
-        MlMemoryTracker memoryTracker
+        MlMemoryTracker memoryTracker,
+        ProjectResolver projectResolver
     ) {
         super(
             MlMemoryAction.NAME,
@@ -75,6 +80,7 @@ public TransportMlMemoryAction(
         );
         this.client = new OriginSettingClient(client, ML_ORIGIN);
         this.memoryTracker = memoryTracker;
+        this.projectResolver = projectResolver;
     }
 
     @Override
@@ -87,6 +93,10 @@ protected void masterOperation(
 
         ClusterSettings clusterSettings = clusterService.getClusterSettings();
 
+        var clusterName = state.getClusterName();
+        var projectMetadata = state.projectState(projectResolver.getProjectId()).metadata();
+        var trainedModelAssignmentMetadata = TrainedModelAssignmentMetadata.fromMetadata(projectMetadata);
+        PersistentTasksCustomMetadata persistentTasksCustomMetadata = projectMetadata.custom(PersistentTasksCustomMetadata.TYPE);
         // Resolve the node specification to some concrete nodes
         String[] nodeIds = state.nodes().resolveNodes(request.getNodeId());
 
@@ -112,7 +122,9 @@ protected void masterOperation(
                         trainedModelCacheInfoRequest,
                         delegate2.delegateFailureAndWrap(
                             (l, trainedModelCacheInfoResponse) -> handleResponses(
-                                state,
+                                clusterName,
+                                persistentTasksCustomMetadata,
+                                trainedModelAssignmentMetadata,
                                 clusterSettings,
                                 nodesStatsResponse,
                                 trainedModelCacheInfoResponse,
@@ -127,15 +139,14 @@ protected void masterOperation(
         if (memoryTracker.isEverRefreshed()) {
             memoryTrackerRefreshListener.onResponse(null);
         } else {
-            memoryTracker.refresh(
-                state.getMetadata().getProject().custom(PersistentTasksCustomMetadata.TYPE),
-                memoryTrackerRefreshListener
-            );
+            memoryTracker.refresh(persistentTasksCustomMetadata, memoryTrackerRefreshListener);
         }
     }
 
     void handleResponses(
-        ClusterState state,
+        ClusterName clusterName,
+        PersistentTasksCustomMetadata persistentTasks,
+        TrainedModelAssignmentMetadata assignmentMetadata,
         ClusterSettings clusterSettings,
         NodesStatsResponse nodesStatsResponse,
         TrainedModelCacheInfoAction.Response trainedModelCacheInfoResponse,
@@ -173,7 +184,8 @@ void handleResponses(
             ByteSizeValue mlNativeInference;
             if (node.getRoles().contains(DiscoveryNodeRole.ML_ROLE)) {
                 NodeLoad nodeLoad = nodeLoadDetector.detectNodeLoad(
-                    state,
+                    persistentTasks,
+                    assignmentMetadata,
                     node,
                     maxOpenJobsPerNode,
                     maxMachineMemoryPercent,
@@ -219,7 +231,7 @@ void handleResponses(
             );
         }
 
-        listener.onResponse(new MlMemoryAction.Response(state.getClusterName(), nodeResponses, failures));
+        listener.onResponse(new MlMemoryAction.Response(clusterName, nodeResponses, failures));
     }
 
     @Override

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/job/NodeLoadDetector.java

@@ -87,7 +87,24 @@ public NodeLoad detectNodeLoad(
         int maxMachineMemoryPercent,
         boolean useAutoMachineMemoryCalculation
     ) {
-        PersistentTasksCustomMetadata persistentTasks = clusterState.getMetadata().getProject().custom(PersistentTasksCustomMetadata.TYPE);
+        return detectNodeLoad(
+            clusterState.getMetadata().getProject().custom(PersistentTasksCustomMetadata.TYPE),
+            assignmentMetadata,
+            node,
+            maxNumberOfOpenJobs,
+            maxMachineMemoryPercent,
+            useAutoMachineMemoryCalculation
+        );
+    }
+
+    public NodeLoad detectNodeLoad(
+        PersistentTasksCustomMetadata persistentTasks,
+        TrainedModelAssignmentMetadata assignmentMetadata,
+        DiscoveryNode node,
+        int maxNumberOfOpenJobs,
+        int maxMachineMemoryPercent,
+        boolean useAutoMachineMemoryCalculation
+    ) {
         Map<String, String> nodeAttributes = node.getAttributes();
         List<String> errors = new ArrayList<>();
         OptionalLong maxMlMemory = NativeMemoryCalculator.allowedBytesForMl(node, maxMachineMemoryPercent, useAutoMachineMemoryCalculation);