[ML] Make AssignmentPlan to consider only assigned allocations #132170

valeriy42 · 2025-07-30T10:45:20Z

A follow-up to #131990. This PR ensures that only assigned allocations and not current allocations are used in the memory requirements calculation in AssignmentPlan.

This change led to the simplification of the code in ZoneAwareAssignmentPlanner and TrainedModelRebalancer.

This PR also improves readability by adding comments, code documentation, renaming variables, and making the flow of if statements more straightforward.

Marking is a non-issue since the bug was already documented in #131990.

…ignmentPlannerTests

elasticsearchmachine · 2025-07-30T10:45:44Z

Hi @valeriy42, I've created a changelog YAML for you.

# Conflicts: # x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentRebalancer.java # x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AssignmentPlan.java # x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/ZoneAwareAssignmentPlanner.java

valeriy42 · 2025-07-30T13:44:55Z

...n/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentRebalancer.java

+            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());
+                dest.assignModelToNode(deployment, node, sourceAssignment.getValue());


Now, calling assignedModelToNode is enough to correctly account for allocated memory.

Please add a unit test for the copyAssignments method in TrainedModelAssignmentRebalancerTests

valeriy42 · 2025-07-30T13:45:10Z

...ava/org/elasticsearch/xpack/ml/inference/assignment/planning/ZoneAwareAssignmentPlanner.java

-                    );
-                    planBuilder.accountMemory(m, originalNode, requiredMemory);
-                }
+                finalPlanBuilder.assignModelToNode(originalDeployment, originalNode, assignment.getValue());


Now, calling assignedModelToNode is enough to correctly account for allocated memory.

Please make AssignmentPlan::accountMemory() a private method so it is not accidentally called by one of the planners

Copilot

Pull Request Overview

This PR fixes the AssignmentPlan to consider only assigned allocations (not current allocations) in memory requirement calculations. The main change eliminates the need for manual memory accounting of current allocations since AssignmentPlan.Builder now properly calculates memory requirements based only on newly assigned allocations.

Key changes:

Memory calculation logic updated to exclude current allocations from requirements
Simplified code in ZoneAwareAssignmentPlanner and TrainedModelRebalancer by removing manual memory accounting
Improved code readability with better documentation and variable naming

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
AssignmentPlanTests.java	Updated test expectations to reflect corrected memory calculations
ZoneAwareAssignmentPlanner.java	Removed manual memory accounting for current allocations and improved documentation
AssignmentPlanner.java	Enhanced documentation and improved code flow with clearer variable names
AssignmentPlan.java	Added comprehensive documentation and simplified allocation tracking methods
TrainedModelAssignmentRebalancer.java	Removed manual memory accounting logic and improved variable naming

...rc/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AssignmentPlanner.java

...l/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AssignmentPlan.java

elasticsearchmachine · 2025-07-30T13:45:56Z

Pinging @elastic/ml-core (Team:ML)

…rence/assignment/planning/AssignmentPlan.java Co-authored-by: Copilot <[email protected]>

…rence/assignment/planning/AssignmentPlanner.java Co-authored-by: Copilot <[email protected]>

jan-elastic · 2025-07-30T14:20:26Z

.../test/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AssignmentPlanTests.java

@@ -131,7 +132,7 @@ public void testAssignModelToNode_GivenNewPlanSatisfiesCurrentAssignment() {
            builder.assignModelToNode(m, n, 1);

            assertThat(builder.getRemainingCores(n), equalTo(2));
-            assertThat(builder.getRemainingMemory(n), equalTo(ByteSizeValue.ofMb(350).getBytes()));
+            assertThat(builder.getRemainingMemory(n), equalTo(ByteSizeValue.ofMb(50).getBytes()));


Why 50?

If I understand correctly, the node has 350MB and you've assigned a 30MB deployment to it?

I see it consists of:

240MB of memory overhead

2 x 30MB (=model size)

These tests could use some explanation of the numbers.

jan-elastic

LGTM, see the one comment.

(Disclaimer: this code is pretty hard to read, so not 100% sure. Another pair of eyes might be worth it...)

elasticsearchmachine · 2025-07-31T14:26:52Z

💔 Backport failed

Status	Branch	Result
✅	9.1
❌	8.19	Commit could not be cherrypicked due to conflicts
✅	9.0
❌	8.18	Commit could not be cherrypicked due to conflicts

You can use sqren/backport to manually backport by running backport --upstream elastic/elasticsearch --pr 132170

…ic#132170) A follow-up to elastic#131990. This PR ensures that only assigned allocations and not current allocations are used in the memory requirements calculation in AssignmentPlan. This change led to the simplification of the code in ZoneAwareAssignmentPlanner and TrainedModelRebalancer. This PR also improves readability by adding comments, code documentation, renaming variables, and making the flow of if statements more straightforward. Marking is a non-issue since the bug was already documented in elastic#131990.

valeriy42 · 2025-07-31T14:53:58Z

💚 All backports created successfully

Status	Branch	Result
✅	8.19
✅	8.18

Questions ?

Please refer to the Backport tool documentation

…ic#132170) A follow-up to elastic#131990. This PR ensures that only assigned allocations and not current allocations are used in the memory requirements calculation in AssignmentPlan. This change led to the simplification of the code in ZoneAwareAssignmentPlanner and TrainedModelRebalancer. This PR also improves readability by adding comments, code documentation, renaming variables, and making the flow of if statements more straightforward. Marking is a non-issue since the bug was already documented in elastic#131990. (cherry picked from commit 80c47f3)

…) (#132271) A follow-up to #131990. This PR ensures that only assigned allocations and not current allocations are used in the memory requirements calculation in AssignmentPlan. This change led to the simplification of the code in ZoneAwareAssignmentPlanner and TrainedModelRebalancer. This PR also improves readability by adding comments, code documentation, renaming variables, and making the flow of if statements more straightforward. Marking is a non-issue since the bug was already documented in #131990.

…) (#132270) A follow-up to #131990. This PR ensures that only assigned allocations and not current allocations are used in the memory requirements calculation in AssignmentPlan. This change led to the simplification of the code in ZoneAwareAssignmentPlanner and TrainedModelRebalancer. This PR also improves readability by adding comments, code documentation, renaming variables, and making the flow of if statements more straightforward. Marking is a non-issue since the bug was already documented in #131990.

…ic#132170) A follow-up to elastic#131990. This PR ensures that only assigned allocations and not current allocations are used in the memory requirements calculation in AssignmentPlan. This change led to the simplification of the code in ZoneAwareAssignmentPlanner and TrainedModelRebalancer. This PR also improves readability by adding comments, code documentation, renaming variables, and making the flow of if statements more straightforward. Marking is a non-issue since the bug was already documented in elastic#131990.

…#132170) (#132275) * [ML] Make AssignmentPlan to consider only assigned allocations (#132170) A follow-up to #131990. This PR ensures that only assigned allocations and not current allocations are used in the memory requirements calculation in AssignmentPlan. This change led to the simplification of the code in ZoneAwareAssignmentPlanner and TrainedModelRebalancer. This PR also improves readability by adding comments, code documentation, renaming variables, and making the flow of if statements more straightforward. Marking is a non-issue since the bug was already documented in #131990. (cherry picked from commit 80c47f3) * Fix backport errors * Fix unit test

valeriy42 added 3 commits July 30, 2025 11:18

use only assigned allocations

6a9779e

fixed AssignmentPlanTests

a8cc385

fixed tests in TrainedModelAssignmentRebalancerTests and ZoneAwareAss…

d537939

…ignmentPlannerTests

valeriy42 added >bug :ml Machine learning v9.2.0 labels Jul 30, 2025

valeriy42 and others added 4 commits July 30, 2025 12:45

Update docs/changelog/132170.yaml

914381c

[CI] Auto commit changes from spotless

0221a2b

add comments

9445c01

valeriy42 added >non-issue and removed >bug labels Jul 30, 2025

Delete docs/changelog/132170.yaml

cd87fc9

valeriy42 added v9.1.0 v8.19.1 v9.0.5 v8.18.5 auto-backport Automatically create backport pull requests when merged labels Jul 30, 2025

valeriy42 changed the title ~~[ML] Fix double-accounting for allocations~~ [ML] Make AssignmentPlan to consider only assigned allocations Jul 30, 2025

valeriy42 commented Jul 30, 2025

View reviewed changes

valeriy42 requested a review from Copilot July 30, 2025 13:45

valeriy42 marked this pull request as ready for review July 30, 2025 13:45

valeriy42 requested a review from jan-elastic July 30, 2025 13:45

Copilot AI reviewed Jul 30, 2025

View reviewed changes

...rc/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AssignmentPlanner.java Outdated Show resolved Hide resolved

...l/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AssignmentPlan.java Outdated Show resolved Hide resolved

elasticsearchmachine added the Team:ML Meta label for the ML team label Jul 30, 2025

valeriy42 and others added 2 commits July 30, 2025 15:52

Update x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/infe…

b46e15e

…rence/assignment/planning/AssignmentPlan.java Co-authored-by: Copilot <[email protected]>

Update x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/infe…

a48862e

…rence/assignment/planning/AssignmentPlanner.java Co-authored-by: Copilot <[email protected]>

jan-elastic reviewed Jul 30, 2025

View reviewed changes

jan-elastic approved these changes Jul 30, 2025

View reviewed changes

valeriy42 requested a review from davidkyle July 30, 2025 14:32

valeriy42 added 4 commits July 31, 2025 09:44

Merge branch 'main' into bug/allocations-5

4626d87

clean up code

a2247bb

Merge branch 'main' into bug/allocations-5

9521646

Added unit test for TrainedModelAssignmentRebalancer::copyAssignments

007afaa

valeriy42 merged commit 80c47f3 into elastic:main Jul 31, 2025
33 checks passed

valeriy42 deleted the bug/allocations-5 branch July 31, 2025 14:25

This was referenced Jul 31, 2025

[9.1] [ML] Make AssignmentPlan to consider only assigned allocations (#132170) #132270

Merged

[9.0] [ML] Make AssignmentPlan to consider only assigned allocations (#132170) #132271

Merged

elasticsearchmachine added the backport pending label Jul 31, 2025

This was referenced Jul 31, 2025

[8.19] [ML] Make AssignmentPlan to consider only assigned allocations (#132170) #132274

Open

[8.18] [ML] Make AssignmentPlan to consider only assigned allocations (#132170) #132275

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[ML] Make AssignmentPlan to consider only assigned allocations #132170

[ML] Make AssignmentPlan to consider only assigned allocations #132170

Uh oh!

valeriy42 commented Jul 30, 2025 •

edited

Loading

Uh oh!

elasticsearchmachine commented Jul 30, 2025

Uh oh!

valeriy42 Jul 30, 2025

Uh oh!

davidkyle Jul 31, 2025

Uh oh!

valeriy42 Jul 30, 2025

Uh oh!

davidkyle Jul 31, 2025

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

elasticsearchmachine commented Jul 30, 2025

Uh oh!

jan-elastic Jul 30, 2025

Uh oh!

jan-elastic Jul 30, 2025

Uh oh!

jan-elastic left a comment

Uh oh!

Uh oh!

elasticsearchmachine commented Jul 31, 2025

Uh oh!

valeriy42 commented Jul 31, 2025

Uh oh!

Uh oh!

[ML] Make AssignmentPlan to consider only assigned allocations #132170

[ML] Make AssignmentPlan to consider only assigned allocations #132170

Uh oh!

Conversation

valeriy42 commented Jul 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

elasticsearchmachine commented Jul 30, 2025

Uh oh!

valeriy42 Jul 30, 2025

Choose a reason for hiding this comment

Uh oh!

davidkyle Jul 31, 2025

Choose a reason for hiding this comment

Uh oh!

valeriy42 Jul 30, 2025

Choose a reason for hiding this comment

Uh oh!

davidkyle Jul 31, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

elasticsearchmachine commented Jul 30, 2025

Uh oh!

jan-elastic Jul 30, 2025

Choose a reason for hiding this comment

Uh oh!

jan-elastic Jul 30, 2025

Choose a reason for hiding this comment

Uh oh!

jan-elastic left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

elasticsearchmachine commented Jul 31, 2025

💔 Backport failed

Uh oh!

valeriy42 commented Jul 31, 2025

💚 All backports created successfully

Questions ?

Uh oh!

Uh oh!

valeriy42 commented Jul 30, 2025 •

edited

Loading