formatting

Author: valeriy42

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

@@ -138,7 +138,7 @@ static void copyAssignments(AssignmentPlan source, AssignmentPlan.Builder dest,
             Map<AssignmentPlan.Node, Integer> sourceNodeAssignments = source.assignments(deployment).orElse(Map.of());
             for (Map.Entry<AssignmentPlan.Node, Integer> sourceAssignment : sourceNodeAssignments.entrySet()) {
                 AssignmentPlan.Node node = originalNodeById.get(sourceAssignment.getKey().id());
-                if(dest.canAssign(deployment, node, sourceAssignment.getValue())) {
+                if (dest.canAssign(deployment, node, sourceAssignment.getValue())) {
                     dest.assignModelToNode(deployment, node, sourceAssignment.getValue());
                 }
             }
@@ -320,8 +320,10 @@ private Map<List<String>, List<AssignmentPlan.Node>> createNodesByZoneMap() {
     }
 
     private static long getNodeFreeMemoryExcludingPerNodeOverheadAndNativeInference(NodeLoad load) {
-        // load.getFreeMemoryExcludingPerNodeOverhead() = maxMemory - assignedJobMemoryExcludingPerNodeOverhead - 30MB native executable code overhead
-        // assignedJobMemoryExcludingPerNodeOverhead = assignedAnomalyDetectorMemory + assignedDataFrameAnalyticsMemory + assignedNativeInferenceMemory
+        // load.getFreeMemoryExcludingPerNodeOverhead() = maxMemory - assignedJobMemoryExcludingPerNodeOverhead - 30MB native executable
+        // code overhead
+        // assignedJobMemoryExcludingPerNodeOverhead = assignedAnomalyDetectorMemory + assignedDataFrameAnalyticsMemory +
+        // assignedNativeInferenceMemory
         // load.getAssignedNativeInferenceMemory() = assignedNativeInferenceMemory
         // TODO: (valeriy) assignedNativeInferenceMemory is double counted in the current calculation.
         return load.getFreeMemoryExcludingPerNodeOverhead()/* - load.getAssignedNativeInferenceMemory()*/;
@@ -412,17 +414,20 @@ private Optional<String> explainAssignment(
         if (Strings.isNullOrEmpty(load.getError()) == false) {
             return Optional.of(load.getError());
         }
-        // TODO (valeriy): this test should be actually true, but it is false, because we use the "naked" deployment footprint
-        // Get existing allocations for this node to avoid double counting
         int existingAllocationsOnNode = assignmentPlan.assignments(deployment)
-            .flatMap(assignments -> assignments.entrySet().stream()
-                .filter(entry -> entry.getKey().id().equals(node.getId()))
-                .findFirst()
-                .map(Map.Entry::getValue))
+            .flatMap(
+                assignments -> assignments.entrySet()
+                    .stream()
+                    .filter(entry -> entry.getKey().id().equals(node.getId()))
+                    .findFirst()
+                    .map(Map.Entry::getValue)
+            )
             .orElse(0);
         int notYetAssignedAllocations = deployment.allocations() - assignmentPlan.totalAllocations(deployment);
-//        if (deployment.estimateMemoryUsageBytes(deployment.allocations() - existingAllocationsOnNode) > assignmentPlan.getRemainingNodeMemory(node.getId())) {
-        if (deployment.estimateAdditionalMemoryUsageBytes(existingAllocationsOnNode, existingAllocationsOnNode + notYetAssignedAllocations) > assignmentPlan.getRemainingNodeMemory(node.getId())) {
+        if (deployment.estimateAdditionalMemoryUsageBytes(
+            existingAllocationsOnNode,
+            existingAllocationsOnNode + notYetAssignedAllocations
+        ) > assignmentPlan.getRemainingNodeMemory(node.getId())) {
             // If any ML processes are running on a node we require some space to load the shared libraries.
             // So if none are currently running then this per-node overhead must be added to the requirement.
             // From node load we know if we had any jobs or models assigned before the rebalance.

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentRebalancerTests.java

@@ -1268,26 +1268,18 @@ public void testRebalance_GivenDeploymentWithMemoryRequirements_ConsidersNativeE
         long modelMemory = ByteSizeValue.ofMb(200).getBytes();
         long memoryOverhead = ByteSizeValue.ofMb(240).getBytes();
         long JVMOverhead = ByteSizeValue.ofMb(50).getBytes();
-        long nodeMemory = memoryOverhead + modelMemory*2 + JVMOverhead;
+        long nodeMemory = memoryOverhead + modelMemory * 2 + JVMOverhead;
 
         DiscoveryNode node = buildNode("node-1", nodeMemory, 4);
 
         String deploymentId = "model-with-overhead-test";
-        StartTrainedModelDeploymentAction.TaskParams taskParams = normalPriorityParams(
-            deploymentId,
-            deploymentId,
-            modelMemory,
-            1,
-            1
-        );
+        StartTrainedModelDeploymentAction.TaskParams taskParams = normalPriorityParams(deploymentId, deploymentId, modelMemory, 1, 1);
 
         TrainedModelAssignmentMetadata currentMetadata = TrainedModelAssignmentMetadata.Builder.empty().build();
         Map<DiscoveryNode, NodeLoad> nodeLoads = new HashMap<>();
 
         // This node has no jobs or models yet, so the overhead should be accounted for
-        nodeLoads.put(node, NodeLoad.builder("node-1")
-            .setMaxMemory(nodeMemory)
-            .build());
+        nodeLoads.put(node, NodeLoad.builder("node-1").setMaxMemory(nodeMemory).build());
 
         TrainedModelAssignmentMetadata result = new TrainedModelAssignmentRebalancer(
             currentMetadata,
@@ -1310,18 +1302,17 @@ public void testRebalance_GivenDeploymentWithMemoryRequirements_ConsidersNativeE
         DiscoveryNode insufficientNode = buildNode("node-2", insufficientNodeMemory, 4);
 
         Map<DiscoveryNode, NodeLoad> insufficientNodeLoads = Map.of(
-            insufficientNode, NodeLoad.builder("node-2")
-                .setMaxMemory(insufficientNodeMemory)
-                .build()
+            insufficientNode,
+            NodeLoad.builder("node-2").setMaxMemory(insufficientNodeMemory).build()
         );
 
         TrainedModelAssignmentMetadata insufficientResult = new TrainedModelAssignmentRebalancer(
             TrainedModelAssignmentMetadata.Builder.empty().build(),
             insufficientNodeLoads,
             Map.of(List.of(), List.of(insufficientNode)),
             Optional.of(new CreateTrainedModelAssignmentAction.Request(taskParams, null)),
-            1)
-            .rebalance().build();
+            1
+        ).rebalance().build();
 
         TrainedModelAssignment insufficientAssignment = insufficientResult.getDeploymentAssignment(deploymentId);
         assertThat(insufficientAssignment, is(notNullValue()));