Skip to content

Adaptive allocations improvements #126307

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Apr 7, 2025
Merged

Adaptive allocations improvements #126307

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e430d32
Adaptive allocations: don't update deployments that aren't started.
jan-elastic Apr 4, 2025
8bcb917
AssignmentPlanner: don't plan deployments with zero allocations
jan-elastic Apr 4, 2025
b43d2e6
Update JavaDoc
jan-elastic Apr 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ...ain/java/org/elasticsearch/xpack/core/ml/inference/assignment/TrainedModelAssignment.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@ private TrainedModelAssignment(
* @param assignmentState used to track the state of the assignment for rebalancing, autoscaling, and more
* @param reason may contain a human-readable explanation for the current state
* @param startTime the time when the assignment was created
* @param maxAssignedAllocations used for adaptive allocations
* @param maxAssignedAllocations keeps track of the maximum number of allocations used for this assignment
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know what this is exactly used for. It looks like only for recording, not for decision making. Definitely not used by adaptive allocations.

* @param adaptiveAllocationsSettings how the assignment should scale based on usage
*/
TrainedModelAssignment(
Expand Down
11 changes: 11 additions & 0 deletions ...lasticsearch/xpack/ml/inference/adaptiveallocations/AdaptiveAllocationsScalerService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import org.elasticsearch.xpack.core.ml.action.CreateTrainedModelAssignmentAction;
import org.elasticsearch.xpack.core.ml.action.GetDeploymentStatsAction;
import org.elasticsearch.xpack.core.ml.action.UpdateTrainedModelDeploymentAction;
import org.elasticsearch.xpack.core.ml.inference.assignment.AssignmentState;
import org.elasticsearch.xpack.core.ml.inference.assignment.AssignmentStats;
import org.elasticsearch.xpack.core.ml.inference.assignment.RoutingState;
import org.elasticsearch.xpack.core.ml.inference.assignment.TrainedModelAssignment;
Expand Down Expand Up @@ -390,13 +391,15 @@ private void processDeploymentStats(GetDeploymentStatsAction.Response statsRespo

Map<String, Stats> recentStatsByDeployment = new HashMap<>();
Map<String, Integer> numberOfAllocations = new HashMap<>();
Map<String, AssignmentState> assignmentStates = new HashMap<>();
// Check for recent scale ups in the deployment stats, because a different node may have
// caused a scale up when an inference request arrives and there were zero allocations.
Set<String> hasRecentObservedScaleUp = new HashSet<>();

for (AssignmentStats assignmentStats : statsResponse.getStats().results()) {
String deploymentId = assignmentStats.getDeploymentId();
numberOfAllocations.put(deploymentId, assignmentStats.getNumberOfAllocations());
assignmentStates.put(deploymentId, assignmentStats.getState());
Map<String, Stats> deploymentStats = lastInferenceStatsByDeploymentAndNode.computeIfAbsent(
deploymentId,
key -> new HashMap<>()
Expand Down Expand Up @@ -449,6 +452,14 @@ private void processDeploymentStats(GetDeploymentStatsAction.Response statsRespo
logger.debug("adaptive allocations scaler: skipping scaling down [{}] because of recent scaleup.", deploymentId);
continue;
}
if (assignmentStates.get(deploymentId) != AssignmentState.STARTED) {
logger.debug(
"adaptive allocations scaler: skipping scaling [{}] because it is in [{}] state.",
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updating a model that's not in the STARTED state leads to errors in the logs.

deploymentId,
assignmentStates.get(deploymentId)
);
continue;
}
if (newNumberOfAllocations > numberOfAllocations.get(deploymentId)) {
lastScaleUpTimesMillis.put(deploymentId, now);
}
Expand Down
5 changes: 4 additions & 1 deletion ...main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AssignmentPlanner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,10 @@ public class AssignmentPlanner {

public AssignmentPlanner(List<Node> nodes, List<AssignmentPlan.Deployment> deployments) {
this.nodes = nodes.stream().sorted(Comparator.comparing(Node::id)).toList();
this.deployments = deployments.stream().sorted(Comparator.comparing(AssignmentPlan.Deployment::deploymentId)).toList();
this.deployments = deployments.stream()
.filter(deployment -> deployment.allocations() > 0)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If these aren't filtered, computePlan will

Plan with at least one allocation for previously assigned models

.sorted(Comparator.comparing(AssignmentPlan.Deployment::deploymentId))
.toList();
}

public AssignmentPlan computePlan() {
Expand Down
90 changes: 74 additions & 16 deletions ...csearch/xpack/ml/inference/adaptiveallocations/AdaptiveAllocationsScalerServiceTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import org.elasticsearch.xpack.core.ml.action.StartTrainedModelDeploymentAction;
import org.elasticsearch.xpack.core.ml.action.UpdateTrainedModelDeploymentAction;
import org.elasticsearch.xpack.core.ml.inference.assignment.AdaptiveAllocationsSettings;
import org.elasticsearch.xpack.core.ml.inference.assignment.AssignmentState;
import org.elasticsearch.xpack.core.ml.inference.assignment.AssignmentStats;
import org.elasticsearch.xpack.core.ml.inference.assignment.Priority;
import org.elasticsearch.xpack.core.ml.inference.assignment.TrainedModelAssignment;
Expand Down Expand Up @@ -85,7 +86,7 @@ public void tearDown() throws Exception {
super.tearDown();
}

private ClusterState getClusterState(int numAllocations) {
private ClusterState getClusterState(int numAllocations, AssignmentState assignmentState) {
ClusterState clusterState = mock(ClusterState.class);
Metadata metadata = mock(Metadata.class);
when(clusterState.metadata()).thenReturn(metadata);
Expand All @@ -107,7 +108,7 @@ private ClusterState getClusterState(int numAllocations) {
100_000_000
),
new AdaptiveAllocationsSettings(true, null, null)
).build()
).setAssignmentState(assignmentState).build()
)
)
);
Expand All @@ -118,7 +119,8 @@ private GetDeploymentStatsAction.Response getDeploymentStatsResponse(
int numAllocations,
int inferenceCount,
double latency,
boolean recentStartup
boolean recentStartup,
AssignmentState assignmentState
) {
return new GetDeploymentStatsAction.Response(
List.of(),
Expand Down Expand Up @@ -155,15 +157,15 @@ private GetDeploymentStatsAction.Response getDeploymentStatsResponse(
)
),
Priority.NORMAL
)
).setState(assignmentState)
),
0
);
}

public void test_scaleUp() {
// Initialize the cluster with a deployment with 1 allocation.
ClusterState clusterState = getClusterState(1);
ClusterState clusterState = getClusterState(1, AssignmentState.STARTED);
when(clusterService.state()).thenReturn(clusterState);

AdaptiveAllocationsScalerService service = new AdaptiveAllocationsScalerService(
Expand All @@ -189,7 +191,7 @@ public void test_scaleUp() {
doAnswer(invocationOnMock -> {
@SuppressWarnings("unchecked")
var listener = (ActionListener<GetDeploymentStatsAction.Response>) invocationOnMock.getArguments()[2];
listener.onResponse(getDeploymentStatsResponse(1, 1, 11.0, false));
listener.onResponse(getDeploymentStatsResponse(1, 1, 11.0, false, AssignmentState.STARTED));
return Void.TYPE;
}).when(client).execute(eq(GetDeploymentStatsAction.INSTANCE), eq(new GetDeploymentStatsAction.Request("test-deployment")), any());

Expand All @@ -205,7 +207,7 @@ public void test_scaleUp() {
doAnswer(invocationOnMock -> {
@SuppressWarnings("unchecked")
var listener = (ActionListener<GetDeploymentStatsAction.Response>) invocationOnMock.getArguments()[2];
listener.onResponse(getDeploymentStatsResponse(1, 150, 10.0, false));
listener.onResponse(getDeploymentStatsResponse(1, 150, 10.0, false, AssignmentState.STARTED));
return Void.TYPE;
}).when(client).execute(eq(GetDeploymentStatsAction.INSTANCE), eq(new GetDeploymentStatsAction.Request("test-deployment")), any());
doAnswer(invocationOnMock -> {
Expand All @@ -226,7 +228,7 @@ public void test_scaleUp() {
verifyNoMoreInteractions(client, clusterService);
reset(client, clusterService);

clusterState = getClusterState(2);
clusterState = getClusterState(2, AssignmentState.STARTED);
ClusterChangedEvent clusterChangedEvent = mock(ClusterChangedEvent.class);
when(clusterChangedEvent.state()).thenReturn(clusterState);
service.clusterChanged(clusterChangedEvent);
Expand All @@ -236,7 +238,7 @@ public void test_scaleUp() {
doAnswer(invocationOnMock -> {
@SuppressWarnings("unchecked")
var listener = (ActionListener<GetDeploymentStatsAction.Response>) invocationOnMock.getArguments()[2];
listener.onResponse(getDeploymentStatsResponse(2, 0, 9.0, false));
listener.onResponse(getDeploymentStatsResponse(2, 0, 9.0, false, AssignmentState.STARTED));
return Void.TYPE;
}).when(client).execute(eq(GetDeploymentStatsAction.INSTANCE), eq(new GetDeploymentStatsAction.Request("test-deployment")), any());
doAnswer(invocationOnMock -> {
Expand All @@ -257,7 +259,7 @@ public void test_scaleUp() {

public void test_scaleDownToZero_whenNoRequests() {
// Initialize the cluster with a deployment with 1 allocation.
ClusterState clusterState = getClusterState(1);
ClusterState clusterState = getClusterState(1, AssignmentState.STARTED);
when(clusterService.state()).thenReturn(clusterState);

AdaptiveAllocationsScalerService service = new AdaptiveAllocationsScalerService(
Expand All @@ -283,7 +285,7 @@ public void test_scaleDownToZero_whenNoRequests() {
doAnswer(invocationOnMock -> {
@SuppressWarnings("unchecked")
var listener = (ActionListener<GetDeploymentStatsAction.Response>) invocationOnMock.getArguments()[2];
listener.onResponse(getDeploymentStatsResponse(1, 1, 11.0, false));
listener.onResponse(getDeploymentStatsResponse(1, 1, 11.0, false, AssignmentState.STARTED));
return Void.TYPE;
}).when(client).execute(eq(GetDeploymentStatsAction.INSTANCE), eq(new GetDeploymentStatsAction.Request("test-deployment")), any());

Expand All @@ -299,7 +301,7 @@ public void test_scaleDownToZero_whenNoRequests() {
doAnswer(invocationOnMock -> {
@SuppressWarnings("unchecked")
var listener = (ActionListener<GetDeploymentStatsAction.Response>) invocationOnMock.getArguments()[2];
listener.onResponse(getDeploymentStatsResponse(1, 0, 10.0, false));
listener.onResponse(getDeploymentStatsResponse(1, 0, 10.0, false, AssignmentState.STARTED));
return Void.TYPE;
}).when(client).execute(eq(GetDeploymentStatsAction.INSTANCE), eq(new GetDeploymentStatsAction.Request("test-deployment")), any());
doAnswer(invocationOnMock -> {
Expand All @@ -322,9 +324,65 @@ public void test_scaleDownToZero_whenNoRequests() {
service.stop();
}

public void test_dontScale_whenNotStarted() {
// Initialize the cluster with a deployment with 1 allocation.
ClusterState clusterState = getClusterState(1, AssignmentState.STARTING);
when(clusterService.state()).thenReturn(clusterState);

AdaptiveAllocationsScalerService service = new AdaptiveAllocationsScalerService(
threadPool,
clusterService,
client,
inferenceAuditor,
meterRegistry,
true,
1,
1,
2_000
);
service.start();

verify(clusterService).state();
verify(clusterService).addListener(same(service));
verifyNoMoreInteractions(client, clusterService);
reset(client, clusterService);

// First cycle: many inference requests
when(client.threadPool()).thenReturn(threadPool);
doAnswer(invocationOnMock -> {
@SuppressWarnings("unchecked")
var listener = (ActionListener<GetDeploymentStatsAction.Response>) invocationOnMock.getArguments()[2];
listener.onResponse(getDeploymentStatsResponse(1, 10000, 10.0, false, AssignmentState.STARTING));
return Void.TYPE;
}).when(client).execute(eq(GetDeploymentStatsAction.INSTANCE), eq(new GetDeploymentStatsAction.Request("test-deployment")), any());

safeSleep(1200);

verify(client, times(1)).threadPool();
verify(client, times(1)).execute(eq(GetDeploymentStatsAction.INSTANCE), any(), any());
verifyNoMoreInteractions(client, clusterService);
reset(client, clusterService);

// Second cycle: again many inference requests
when(client.threadPool()).thenReturn(threadPool);
doAnswer(invocationOnMock -> {
@SuppressWarnings("unchecked")
var listener = (ActionListener<GetDeploymentStatsAction.Response>) invocationOnMock.getArguments()[2];
listener.onResponse(getDeploymentStatsResponse(1, 20000, 10.0, false, AssignmentState.STARTING));
return Void.TYPE;
}).when(client).execute(eq(GetDeploymentStatsAction.INSTANCE), eq(new GetDeploymentStatsAction.Request("test-deployment")), any());

safeSleep(1200);

verify(client, times(1)).threadPool();
verify(client, times(1)).execute(eq(GetDeploymentStatsAction.INSTANCE), any(), any());
verifyNoMoreInteractions(client, clusterService);
service.stop();
}

public void test_noScaleDownToZero_whenRecentlyScaledUpByOtherNode() {
// Initialize the cluster with a deployment with 1 allocation.
ClusterState clusterState = getClusterState(1);
ClusterState clusterState = getClusterState(1, AssignmentState.STARTED);
when(clusterService.state()).thenReturn(clusterState);

AdaptiveAllocationsScalerService service = new AdaptiveAllocationsScalerService(
Expand All @@ -350,7 +408,7 @@ public void test_noScaleDownToZero_whenRecentlyScaledUpByOtherNode() {
doAnswer(invocationOnMock -> {
@SuppressWarnings("unchecked")
var listener = (ActionListener<GetDeploymentStatsAction.Response>) invocationOnMock.getArguments()[2];
listener.onResponse(getDeploymentStatsResponse(1, 1, 11.0, true));
listener.onResponse(getDeploymentStatsResponse(1, 1, 11.0, true, AssignmentState.STARTED));
return Void.TYPE;
}).when(client).execute(eq(GetDeploymentStatsAction.INSTANCE), eq(new GetDeploymentStatsAction.Request("test-deployment")), any());

Expand All @@ -366,7 +424,7 @@ public void test_noScaleDownToZero_whenRecentlyScaledUpByOtherNode() {
doAnswer(invocationOnMock -> {
@SuppressWarnings("unchecked")
var listener = (ActionListener<GetDeploymentStatsAction.Response>) invocationOnMock.getArguments()[2];
listener.onResponse(getDeploymentStatsResponse(1, 0, 10.0, true));
listener.onResponse(getDeploymentStatsResponse(1, 0, 10.0, true, AssignmentState.STARTED));
return Void.TYPE;
}).when(client).execute(eq(GetDeploymentStatsAction.INSTANCE), eq(new GetDeploymentStatsAction.Request("test-deployment")), any());
doAnswer(invocationOnMock -> {
Expand All @@ -388,7 +446,7 @@ public void test_noScaleDownToZero_whenRecentlyScaledUpByOtherNode() {
doAnswer(invocationOnMock -> {
@SuppressWarnings("unchecked")
var listener = (ActionListener<GetDeploymentStatsAction.Response>) invocationOnMock.getArguments()[2];
listener.onResponse(getDeploymentStatsResponse(1, 0, 10.0, false));
listener.onResponse(getDeploymentStatsResponse(1, 0, 10.0, false, AssignmentState.STARTED));
return Void.TYPE;
}).when(client).execute(eq(GetDeploymentStatsAction.INSTANCE), eq(new GetDeploymentStatsAction.Request("test-deployment")), any());
doAnswer(invocationOnMock -> {
Expand Down