Skip to content

Commit 1baedcd

Browse files
[ML] Adding missing onFailure call for Inference API start model request (elastic#126930)
* Adding missing onFailure call * Update docs/changelog/126930.yaml
1 parent 5e2293d commit 1baedcd

File tree

4 files changed

+55
-1
lines changed

4 files changed

+55
-1
lines changed

docs/changelog/126930.yaml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 126930
2+
summary: Adding missing `onFailure` call for Inference API start model request
3+
area: Machine Learning
4+
type: bug
5+
issues: []

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/BaseElasticsearchInternalService.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ public void start(Model model, TimeValue timeout, ActionListener<Boolean> finalL
107107
})
108108
.<Boolean>andThen((l2, modelDidPut) -> {
109109
var startRequest = esModel.getStartTrainedModelDeploymentActionRequest(timeout);
110-
var responseListener = esModel.getCreateTrainedModelAssignmentActionListener(model, finalListener);
110+
var responseListener = esModel.getCreateTrainedModelAssignmentActionListener(model, l2);
111111
client.execute(StartTrainedModelDeploymentAction.INSTANCE, startRequest, responseListener);
112112
})
113113
.addListener(finalListener);

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalModel.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,8 @@ public void onFailure(Exception e) {
105105
&& statusException.getRootCause() instanceof ResourceAlreadyExistsException) {
106106
// Deployment is already started
107107
listener.onResponse(Boolean.TRUE);
108+
} else {
109+
listener.onFailure(e);
108110
}
109111
return;
110112
}

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElasticsearchInternalServiceTests.java

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import org.elasticsearch.inference.ModelConfigurations;
3737
import org.elasticsearch.inference.SimilarityMeasure;
3838
import org.elasticsearch.inference.TaskType;
39+
import org.elasticsearch.rest.RestStatus;
3940
import org.elasticsearch.test.ESTestCase;
4041
import org.elasticsearch.threadpool.ThreadPool;
4142
import org.elasticsearch.xcontent.ParseField;
@@ -47,13 +48,16 @@
4748
import org.elasticsearch.xpack.core.inference.results.ChunkedInferenceEmbeddingSparse;
4849
import org.elasticsearch.xpack.core.inference.results.ChunkedInferenceError;
4950
import org.elasticsearch.xpack.core.ml.MachineLearningField;
51+
import org.elasticsearch.xpack.core.ml.action.CreateTrainedModelAssignmentAction;
5052
import org.elasticsearch.xpack.core.ml.action.GetDeploymentStatsAction;
5153
import org.elasticsearch.xpack.core.ml.action.GetTrainedModelsAction;
5254
import org.elasticsearch.xpack.core.ml.action.InferModelAction;
5355
import org.elasticsearch.xpack.core.ml.action.InferTrainedModelDeploymentAction;
5456
import org.elasticsearch.xpack.core.ml.action.PutTrainedModelAction;
57+
import org.elasticsearch.xpack.core.ml.action.StartTrainedModelDeploymentAction;
5558
import org.elasticsearch.xpack.core.ml.inference.TrainedModelConfig;
5659
import org.elasticsearch.xpack.core.ml.inference.TrainedModelPrefixStrings;
60+
import org.elasticsearch.xpack.core.ml.inference.assignment.AdaptiveAllocationsSettings;
5761
import org.elasticsearch.xpack.core.ml.inference.assignment.AssignmentStats;
5862
import org.elasticsearch.xpack.core.ml.inference.results.ErrorInferenceResults;
5963
import org.elasticsearch.xpack.core.ml.inference.results.MlTextEmbeddingResults;
@@ -1870,6 +1874,49 @@ public void testUpdateWithMlEnabled() throws IOException, InterruptedException {
18701874
}
18711875
}
18721876

1877+
public void testStart_OnFailure_WhenTimeoutOccurs() throws IOException {
1878+
var model = new ElserInternalModel(
1879+
"inference_id",
1880+
TaskType.SPARSE_EMBEDDING,
1881+
"elasticsearch",
1882+
new ElserInternalServiceSettings(
1883+
new ElasticsearchInternalServiceSettings(1, 1, "id", new AdaptiveAllocationsSettings(false, 0, 0), null)
1884+
),
1885+
new ElserMlNodeTaskSettings(),
1886+
null
1887+
);
1888+
1889+
var client = mock(Client.class);
1890+
when(client.threadPool()).thenReturn(threadPool);
1891+
1892+
doAnswer(invocationOnMock -> {
1893+
ActionListener<GetTrainedModelsAction.Response> listener = invocationOnMock.getArgument(2);
1894+
var builder = GetTrainedModelsAction.Response.builder();
1895+
builder.setModels(List.of(mock(TrainedModelConfig.class)));
1896+
builder.setTotalCount(1);
1897+
1898+
listener.onResponse(builder.build());
1899+
return Void.TYPE;
1900+
}).when(client).execute(eq(GetTrainedModelsAction.INSTANCE), any(), any());
1901+
1902+
doAnswer(invocationOnMock -> {
1903+
ActionListener<CreateTrainedModelAssignmentAction.Response> listener = invocationOnMock.getArgument(2);
1904+
listener.onFailure(new ElasticsearchStatusException("failed", RestStatus.GATEWAY_TIMEOUT));
1905+
return Void.TYPE;
1906+
}).when(client).execute(eq(StartTrainedModelDeploymentAction.INSTANCE), any(), any());
1907+
1908+
try (var service = createService(client)) {
1909+
var actionListener = new PlainActionFuture<Boolean>();
1910+
service.start(model, TimeValue.timeValueSeconds(30), actionListener);
1911+
var exception = expectThrows(
1912+
ElasticsearchStatusException.class,
1913+
() -> actionListener.actionGet(TimeValue.timeValueSeconds(30))
1914+
);
1915+
1916+
assertThat(exception.getMessage(), is("failed"));
1917+
}
1918+
}
1919+
18731920
private ElasticsearchInternalService createService(Client client) {
18741921
var cs = mock(ClusterService.class);
18751922
var cSettings = new ClusterSettings(Settings.EMPTY, Set.of(MachineLearningField.MAX_LAZY_ML_NODES));

0 commit comments

Comments
 (0)