[ML] Fix test failure updating model deployment with stale cluster state. #128667
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -11,8 +11,6 @@ | |
| import org.apache.logging.log4j.Logger; | ||
| import org.elasticsearch.ElasticsearchStatusException; | ||
| import org.elasticsearch.ResourceNotFoundException; | ||
| import org.elasticsearch.action.ActionListener; | ||
| import org.elasticsearch.action.ActionRequestValidationException; | ||
| import org.elasticsearch.action.support.master.AcknowledgedResponse; | ||
|
|
@@ -79,9 +77,6 @@ public class TrainedModelAssignmentClusterService implements ClusterStateListene | |
|
|
||
| private static final Logger logger = LogManager.getLogger(TrainedModelAssignmentClusterService.class); | ||
|
|
||
| private final ClusterService clusterService; | ||
| private final ThreadPool threadPool; | ||
| private final NodeLoadDetector nodeLoadDetector; | ||
|
|
@@ -169,14 +164,6 @@ public void clusterChanged(ClusterChangedEvent event) { | |
| return; | ||
| } | ||
|
|
||
| if (event.nodesAdded()) { | ||
| logMlNodeHeterogeneity(); | ||
| } | ||
|
|
@@ -203,10 +190,6 @@ public void clusterChanged(ClusterChangedEvent event) { | |
| } | ||
| } | ||
|
|
||
| boolean eventStateHasGlobalBlockStateNotRecoveredBlock(ClusterChangedEvent event) { | ||
| return event.state().blocks().hasGlobalBlock(GatewayService.STATE_NOT_RECOVERED_BLOCK); | ||
| } | ||
|
|
@@ -400,18 +383,6 @@ public void createNewModelAssignment( | |
| CreateTrainedModelAssignmentAction.Request request, | ||
| ActionListener<TrainedModelAssignment> listener | ||
| ) { | ||
| if (MlMetadata.getMlMetadata(clusterService.state()).isResetMode()) { | ||
| listener.onFailure( | ||
| new ElasticsearchStatusException( | ||
|
|
@@ -522,13 +493,11 @@ private static ClusterState update(ClusterState currentState, TrainedModelAssign | |
|
|
||
| private static ClusterState forceUpdate(ClusterState currentState, TrainedModelAssignmentMetadata.Builder modelAssignments) { | ||
| logger.debug(() -> format("updated assignments: %s", modelAssignments.build())); | ||
|
|
||
| ProjectMetadata.Builder builder = ProjectMetadata.builder(currentState.metadata().getProject()); | ||
| builder.putCustom(TrainedModelAssignmentMetadata.NAME, modelAssignments.build()) | ||
| .removeCustom(TrainedModelAssignmentMetadata.DEPRECATED_NAME); | ||
|
|
||
| return ClusterState.builder(currentState).putProjectMetadata(builder).build(); | ||
| } | ||
|
|
||
|
|
@@ -844,7 +813,7 @@ private void updateDeployment( | |
| } | ||
| boolean hasUpdates = hasUpdates(numberOfAllocations, adaptiveAllocationsSettingsUpdates, existingAssignment); | ||
| if (hasUpdates == false) { | ||
| logger.debug("no updates to be made for deployment [{}]", deploymentId); | ||
| listener.onResponse(existingAssignment); | ||
| return; | ||
| } | ||
|
|
@@ -858,27 +827,17 @@ private void updateDeployment( | |
| ); | ||
| return; | ||
| } | ||
|
|
||
| ActionListener<TrainedModelAssignmentMetadata.Builder> updatedAssignmentListener = ActionListener.wrap( | ||
| updatedAssignment -> submitUnbatchedTask("update model deployment", new ClusterStateUpdateTask() { | ||
|
There was a problem hiding this comment. This is the fix, here the new assignment state is passed rather than the updated cluster state. There was a problem hiding this comment. What if the That's a similar bug to the existing one, right? But much smaller, because the Should we protect against that? For example, only replace the Maybe that's overkill and too complicated though. WDYT? There was a problem hiding this comment. If we're deciding not to fix this, let's leave a comment about this small issue with this implementation and call it a day There was a problem hiding this comment. Looking at the By comparing the starting state with the latest state before applying the updated There was a problem hiding this comment. OK, cool! I missed that part of the code. Guess this all works then as is |
||
|
|
||
| private volatile boolean isUpdated; | ||
|
|
||
| @Override | ||
| public ClusterState execute(ClusterState currentState) { | ||
| if (areClusterStatesCompatibleForRebalance(clusterState, currentState)) { | ||
| isUpdated = true; | ||
| return update(currentState, updatedAssignment); | ||
| } | ||
| logger.debug(() -> format("[%s] Retrying update as cluster state has been modified", deploymentId)); | ||
| updateDeployment(currentState, deploymentId, numberOfAllocations, adaptiveAllocationsSettings, isInternal, listener); | ||
|
|
@@ -910,7 +869,7 @@ public void clusterStateProcessed(ClusterState oldState, ClusterState newState) | |
| listener::onFailure | ||
| ); | ||
|
|
||
| updateAssignment(clusterState, existingAssignment, numberOfAllocations, adaptiveAllocationsSettings, updatedAssignmentListener); | ||
| } | ||
|
|
||
| static boolean hasUpdates( | ||
|
|
@@ -944,7 +903,7 @@ private void updateAssignment( | |
| TrainedModelAssignment assignment, | ||
| Integer numberOfAllocations, | ||
| AdaptiveAllocationsSettings adaptiveAllocationsSettings, | ||
| ActionListener<TrainedModelAssignmentMetadata.Builder> listener | ||
| ) { | ||
| threadPool.executor(MachineLearning.UTILITY_THREAD_POOL_NAME).execute(() -> { | ||
| if (numberOfAllocations == null || numberOfAllocations == assignment.getTaskParams().getNumberOfAllocations()) { | ||
|
|
@@ -961,21 +920,21 @@ private void updateAndKeepNumberOfAllocations( | |
| ClusterState clusterState, | ||
| TrainedModelAssignment assignment, | ||
| AdaptiveAllocationsSettings adaptiveAllocationsSettings, | ||
| ActionListener<TrainedModelAssignmentMetadata.Builder> listener | ||
| ) { | ||
| TrainedModelAssignment.Builder updatedAssignment = TrainedModelAssignment.Builder.fromAssignment(assignment) | ||
| .setAdaptiveAllocationsSettings(adaptiveAllocationsSettings); | ||
| TrainedModelAssignmentMetadata.Builder builder = TrainedModelAssignmentMetadata.builder(clusterState); | ||
| builder.updateAssignment(assignment.getDeploymentId(), updatedAssignment); | ||
| listener.onResponse(builder); | ||
| } | ||
|
|
||
| private void increaseNumberOfAllocations( | ||
| ClusterState clusterState, | ||
| TrainedModelAssignment assignment, | ||
| int numberOfAllocations, | ||
| AdaptiveAllocationsSettings adaptiveAllocationsSettings, | ||
| ActionListener<TrainedModelAssignmentMetadata.Builder> listener | ||
| ) { | ||
| try { | ||
| TrainedModelAssignment.Builder updatedAssignment = TrainedModelAssignment.Builder.fromAssignment(assignment) | ||
|
|
@@ -995,7 +954,7 @@ private void increaseNumberOfAllocations( | |
| ) | ||
| ); | ||
| } else { | ||
| listener.onResponse(rebalancedMetadata); | ||
| } | ||
| } catch (Exception e) { | ||
| listener.onFailure(e); | ||
|
|
@@ -1007,7 +966,7 @@ private void decreaseNumberOfAllocations( | |
| TrainedModelAssignment assignment, | ||
| int numberOfAllocations, | ||
| AdaptiveAllocationsSettings adaptiveAllocationsSettings, | ||
| ActionListener<TrainedModelAssignmentMetadata.Builder> listener | ||
| ) { | ||
| TrainedModelAssignment.Builder updatedAssignment = numberOfAllocations < assignment.totalTargetAllocations() | ||
| ? new AllocationReducer(assignment, nodeAvailabilityZoneMapper.buildMlNodesByAvailabilityZone(clusterState)).reduceTo( | ||
|
|
@@ -1022,7 +981,7 @@ private void decreaseNumberOfAllocations( | |
| } | ||
| TrainedModelAssignmentMetadata.Builder builder = TrainedModelAssignmentMetadata.builder(clusterState); | ||
| builder.updateAssignment(assignment.getDeploymentId(), updatedAssignment); | ||
| listener.onResponse(builder); | ||
| } | ||
|
|
||
| static ClusterState setToStopping(ClusterState clusterState, String deploymentId, String reason) { | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
These version checks are redundant in 9.0 and 9.1. The 8.x backports will need to keep them however.