Skip to content

[ML] Flag updates from Inference #131725

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

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
+45 −17
Open

[ML] Flag updates from Inference #131725

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
b124b8f
[ML] Use adaptive allocations in test
prwhelan Jul 22, 2025
061cc5e
Merge branch 'main' into block-updates
prwhelan Jul 22, 2025
9ddaec4
Flag when updates come from the inference endpoint
prwhelan Jul 23, 2025
b97a4ca
Merge branch 'main' of github.com:elastic/elasticsearch into block-up…
prwhelan Jul 29, 2025
0107a7c
Change booleans to enum
prwhelan Jul 29, 2025
00f7ba7
Merge branch 'main' into block-updates
prwhelan Jul 30, 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
1 change: 1 addition & 0 deletions server/src/main/java/org/elasticsearch/TransportVersions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -354,6 +354,7 @@ static TransportVersion def(int id) {
public static final TransportVersion RERANK_SNIPPETS = def(9_130_0_00);
public static final TransportVersion PIPELINE_TRACKING_INFO = def(9_131_0_00);
public static final TransportVersion COMPONENT_TEMPLATE_TRACKING_INFO = def(9_132_0_00);
public static final TransportVersion INFERENCE_UPDATE_ML = def(9_133_0_00);

/*
* STOP! READ THIS FIRST! No, really,
Expand Down
45 changes: 34 additions & 11 deletions .../main/java/org/elasticsearch/xpack/core/ml/action/UpdateTrainedModelDeploymentAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,17 @@
import java.io.IOException;
import java.util.Objects;
Copy link
Contributor

Choose a reason for hiding this comment

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

I'm missing a bit of context: why do we need to distinguish between these cases?

Copy link
Contributor

Choose a reason for hiding this comment

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

Is there a corresponding Serverless PR?

Copy link
Member Author

Choose a reason for hiding this comment

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

Yeah, let me ping you with the internal documentation

Copy link
Member Author

Choose a reason for hiding this comment

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

I'm missing a bit of context: why do we need to distinguish between these cases?

We need to allow updates to num_allocations in serverless that originate from the AdaptiveAllocationsScalerService (ADAPTIVE_ALLOCATIONS), but we want to disallow updates from users (API and INFERENCE). The only alternative I thought of was refactoring AdaptiveAllocationsScalerService to update directly rather than through the API, but that felt more intrusive.


import static org.elasticsearch.TransportVersions.INFERENCE_UPDATE_ML;
import static org.elasticsearch.xpack.core.ml.action.StartTrainedModelDeploymentAction.Request.ADAPTIVE_ALLOCATIONS;
import static org.elasticsearch.xpack.core.ml.action.StartTrainedModelDeploymentAction.Request.MODEL_ID;
import static org.elasticsearch.xpack.core.ml.action.StartTrainedModelDeploymentAction.Request.NUMBER_OF_ALLOCATIONS;

public class UpdateTrainedModelDeploymentAction extends ActionType<CreateTrainedModelAssignmentAction.Response> {
public enum Source {
API,
ADAPTIVE_ALLOCATIONS,
INFERENCE
}

public static final UpdateTrainedModelDeploymentAction INSTANCE = new UpdateTrainedModelDeploymentAction();
public static final String NAME = "cluster:admin/xpack/ml/trained_models/deployment/update";
Expand Down Expand Up @@ -73,7 +79,7 @@ public static Request parseRequest(String deploymentId, XContentParser parser) {
private String deploymentId;
private Integer numberOfAllocations;
private AdaptiveAllocationsSettings adaptiveAllocationsSettings;
private boolean isInternal;
private Source source = Source.API;

private Request() {
super(TRAPPY_IMPLICIT_DEFAULT_MASTER_NODE_TIMEOUT, DEFAULT_ACK_TIMEOUT);
Expand All @@ -90,11 +96,17 @@ public Request(StreamInput in) throws IOException {
if (in.getTransportVersion().before(TransportVersions.V_8_16_0)) {
numberOfAllocations = in.readVInt();
adaptiveAllocationsSettings = null;
isInternal = false;
source = Source.API;
} else {
numberOfAllocations = in.readOptionalVInt();
adaptiveAllocationsSettings = in.readOptionalWriteable(AdaptiveAllocationsSettings::new);
isInternal = in.readBoolean();
if (in.getTransportVersion().before(INFERENCE_UPDATE_ML)) {
// we changed over from a boolean to an enum
// when it was a boolean, true came from adaptive allocations and false came from the rest api
source = in.readBoolean() ? Source.ADAPTIVE_ALLOCATIONS : Source.API;
} else {
source = in.readEnum(Source.class);
}
}
}

Expand All @@ -119,11 +131,15 @@ public void setAdaptiveAllocationsSettings(AdaptiveAllocationsSettings adaptiveA
}

public boolean isInternal() {
return isInternal;
return source == Source.INFERENCE || source == Source.ADAPTIVE_ALLOCATIONS;
Copy link
Contributor

Choose a reason for hiding this comment

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

Can you confirm that we do want Source.INFERENCE here for all the usage of isInternal() below?

Copy link
Member Author

Choose a reason for hiding this comment

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

Confirmed! Yeah inference update code previously set isInternal to true (back when the boolean existed)

}

public void setIsInternal(boolean isInternal) {
this.isInternal = isInternal;
public void setSource(Source source) {
this.source = source != null ? source : this.source;
}

public Source getSource() {
return source;
}

public AdaptiveAllocationsSettings getAdaptiveAllocationsSettings() {
Expand All @@ -139,7 +155,14 @@ public void writeTo(StreamOutput out) throws IOException {
} else {
out.writeOptionalVInt(numberOfAllocations);
out.writeOptionalWriteable(adaptiveAllocationsSettings);
out.writeBoolean(isInternal);
if (out.getTransportVersion().before(INFERENCE_UPDATE_ML)) {
// we changed over from a boolean to an enum
// when it was a boolean, true came from adaptive allocations and false came from the rest api
// treat "inference" as if it came from the api
out.writeBoolean(isInternal());
Copy link
Contributor

Choose a reason for hiding this comment

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

Do we need to determine if source == Source.ADAPTIVE_ALLOCATIONS here? Since this will return true for Source.INFERENCE as well?

Copy link
Member Author

Choose a reason for hiding this comment

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

Previously, we set the boolean to true if the source was either from the inference update api or the adaptive allocations autoscaler. out.writeBoolean(isInternal()) preserves this logic (i think). It means the stream reader will think an inference api call is an adaptive allocations api call, but that only affects serverless which is only mixed cluster during a rolling update.

} else {
out.writeEnum(source);
}
}
}

Expand All @@ -161,10 +184,10 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
public ActionRequestValidationException validate() {
ActionRequestValidationException validationException = new ActionRequestValidationException();
if (numberOfAllocations != null) {
if (numberOfAllocations < 0 || (isInternal == false && numberOfAllocations == 0)) {
if (numberOfAllocations < 0 || (isInternal() == false && numberOfAllocations == 0)) {
validationException.addValidationError("[" + NUMBER_OF_ALLOCATIONS + "] must be a positive integer");
}
if (isInternal == false
if (isInternal() == false
&& adaptiveAllocationsSettings != null
&& adaptiveAllocationsSettings.getEnabled() == Boolean.TRUE) {
validationException.addValidationError(
Expand All @@ -183,7 +206,7 @@ public ActionRequestValidationException validate() {

@Override
public int hashCode() {
return Objects.hash(deploymentId, numberOfAllocations, adaptiveAllocationsSettings, isInternal);
return Objects.hash(deploymentId, numberOfAllocations, adaptiveAllocationsSettings, source);
}

@Override
Expand All @@ -198,7 +221,7 @@ public boolean equals(Object obj) {
return Objects.equals(deploymentId, other.deploymentId)
&& Objects.equals(numberOfAllocations, other.numberOfAllocations)
&& Objects.equals(adaptiveAllocationsSettings, other.adaptiveAllocationsSettings)
&& isInternal == other.isInternal;
&& source == other.source;
}

@Override
Expand Down
2 changes: 1 addition & 1 deletion ...in/java/org/elasticsearch/xpack/inference/action/TransportUpdateInferenceModelAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -294,7 +294,7 @@ private void updateInClusterEndpoint(
var updateRequest = new UpdateTrainedModelDeploymentAction.Request(deploymentId);
updateRequest.setNumberOfAllocations(elasticServiceSettings.getNumAllocations());
updateRequest.setAdaptiveAllocationsSettings(elasticServiceSettings.getAdaptiveAllocationsSettings());
updateRequest.setIsInternal(true);
updateRequest.setSource(UpdateTrainedModelDeploymentAction.Source.INFERENCE);

var delegate = listener.<CreateTrainedModelAssignmentAction.Response>delegateFailure((l2, response) -> {
modelRegistry.updateModelTransaction(newModel, existingParsedModel, l2);
Expand Down
2 changes: 1 addition & 1 deletion ...lasticsearch/xpack/ml/inference/adaptiveallocations/AdaptiveAllocationsScalerService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -520,7 +520,7 @@ private void updateNumberOfAllocations(
) {
UpdateTrainedModelDeploymentAction.Request updateRequest = new UpdateTrainedModelDeploymentAction.Request(deploymentId);
updateRequest.setNumberOfAllocations(numberOfAllocations);
updateRequest.setIsInternal(true);
updateRequest.setSource(UpdateTrainedModelDeploymentAction.Source.ADAPTIVE_ALLOCATIONS);
ClientHelper.executeAsyncWithOrigin(
client,
ClientHelper.ML_ORIGIN,
Expand Down
6 changes: 3 additions & 3 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 @@ -229,7 +229,7 @@ public void test_scaleUp() {
verify(client, times(1)).execute(eq(GetDeploymentStatsAction.INSTANCE), any(), any());
var updateRequest = new UpdateTrainedModelDeploymentAction.Request("test-deployment");
updateRequest.setNumberOfAllocations(2);
updateRequest.setIsInternal(true);
updateRequest.setSource(UpdateTrainedModelDeploymentAction.Source.ADAPTIVE_ALLOCATIONS);
verify(client, times(1)).execute(eq(UpdateTrainedModelDeploymentAction.INSTANCE), eq(updateRequest), any());
verifyNoMoreInteractions(client, clusterService);
reset(client, clusterService);
Expand Down Expand Up @@ -323,7 +323,7 @@ public void test_scaleDownToZero_whenNoRequests() {
verify(client, times(1)).execute(eq(GetDeploymentStatsAction.INSTANCE), any(), any());
var updateRequest = new UpdateTrainedModelDeploymentAction.Request("test-deployment");
updateRequest.setNumberOfAllocations(0);
updateRequest.setIsInternal(true);
updateRequest.setSource(UpdateTrainedModelDeploymentAction.Source.ADAPTIVE_ALLOCATIONS);
verify(client, times(1)).execute(eq(UpdateTrainedModelDeploymentAction.INSTANCE), eq(updateRequest), any());
verifyNoMoreInteractions(client, clusterService);

Expand Down Expand Up @@ -468,7 +468,7 @@ public void test_noScaleDownToZero_whenRecentlyScaledUpByOtherNode() {
verify(client, times(1)).execute(eq(GetDeploymentStatsAction.INSTANCE), any(), any());
var updateRequest = new UpdateTrainedModelDeploymentAction.Request("test-deployment");
updateRequest.setNumberOfAllocations(0);
updateRequest.setIsInternal(true);
updateRequest.setSource(UpdateTrainedModelDeploymentAction.Source.ADAPTIVE_ALLOCATIONS);
verify(client, times(1)).execute(eq(UpdateTrainedModelDeploymentAction.INSTANCE), eq(updateRequest), any());
verifyNoMoreInteractions(client, clusterService);

Expand Down
6 changes: 5 additions & 1 deletion ...ugin/src/yamlRestTest/resources/rest-api-spec/test/ml/update_trained_model_deployment.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,5 +6,9 @@
model_id: "missing-model"
body: >
{
"number_of_allocations": 4
"adaptive_allocations": {
"enabled": true,
"min_number_of_allocations": 0,
"max_number_of_allocations": 1
}
}