Skip to content

[ML] Wait for up to 2 seconds for yellow status before starting search #115938

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 7 commits into from
Oct 31, 2024
Merged

[ML] Wait for up to 2 seconds for yellow status before starting search #115938

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
cb43c45
Wait for up to 2 seconds for yellow before starting search
Oct 29, 2024
74d8de7
Update docs/changelog/115938.yaml
maxhniebergall Oct 30, 2024
ec345fa
Merge branch 'main' into fixTrainedModelStatsWait
elasticmachine Oct 30, 2024
1d5380d
Merge branch 'main' into fixTrainedModelStatsWait
elasticmachine Oct 30, 2024
a2e763b
Update docs/changelog/115938.yaml
maxhniebergall Oct 30, 2024
d97e769
improvements from review
Oct 31, 2024
1ba61da
Merge branch 'main' into fixTrainedModelStatsWait
elasticmachine Oct 31, 2024
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
9 changes: 9 additions & 0 deletions docs/changelog/115938.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
pr: 115938
summary: Wait for up to 2 seconds for yellow status before starting search
area: Machine Learning
type: bug
issues:
- 107777
- 105955
- 107815
- 112191
142 changes: 91 additions & 51 deletions .../src/main/java/org/elasticsearch/xpack/ml/inference/persistence/TrainedModelProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
import org.elasticsearch.ResourceNotFoundException;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.DocWriteRequest;
import org.elasticsearch.action.admin.cluster.health.ClusterHealthRequest;
import org.elasticsearch.action.admin.cluster.health.ClusterHealthResponse;
import org.elasticsearch.action.admin.indices.refresh.RefreshAction;
import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
import org.elasticsearch.action.bulk.BulkItemResponse;
Expand All @@ -28,6 +30,7 @@
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.TransportSearchAction;
import org.elasticsearch.action.support.IndicesOptions;
import org.elasticsearch.action.support.SubscribableListener;
import org.elasticsearch.action.support.WriteRequest;
import org.elasticsearch.action.support.broadcast.BroadcastResponse;
import org.elasticsearch.action.support.master.AcknowledgedResponse;
Expand All @@ -43,6 +46,7 @@
import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
import org.elasticsearch.common.xcontent.XContentHelper;
import org.elasticsearch.core.Nullable;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.core.Tuple;
import org.elasticsearch.index.IndexNotFoundException;
import org.elasticsearch.index.engine.VersionConflictEngineException;
Expand Down Expand Up @@ -105,6 +109,7 @@
import java.util.Objects;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.stream.Collectors;

Expand Down Expand Up @@ -1053,63 +1058,98 @@ public void expandIds(
}

public void getInferenceStats(String[] modelIds, @Nullable TaskId parentTaskId, ActionListener<List<InferenceStats>> listener) {
MultiSearchRequest multiSearchRequest = new MultiSearchRequest();
Arrays.stream(modelIds).map(TrainedModelProvider::buildStatsSearchRequest).forEach(multiSearchRequest::add);
if (multiSearchRequest.requests().isEmpty()) {
listener.onResponse(Collections.emptyList());
return;
}
if (parentTaskId != null) {
multiSearchRequest.setParentTask(parentTaskId);
}
executeAsyncWithOrigin(
client.threadPool().getThreadContext(),
ML_ORIGIN,
multiSearchRequest,
ActionListener.<MultiSearchResponse>wrap(responses -> {
List<InferenceStats> allStats = new ArrayList<>(modelIds.length);
int modelIndex = 0;
assert responses.getResponses().length == modelIds.length : "mismatch between search response size and models requested";
for (MultiSearchResponse.Item response : responses.getResponses()) {
if (response.isFailure()) {
if (ExceptionsHelper.unwrapCause(response.getFailure()) instanceof ResourceNotFoundException) {
modelIndex++;
continue;
}

SubscribableListener.<ClusterHealthResponse>newForked((delegate) -> {
// first wait for the index to be available
executeAsyncWithOrigin(
client.threadPool().getThreadContext(),
ML_ORIGIN,
new ClusterHealthRequest(new TimeValue(2, TimeUnit.SECONDS), MlStatsIndex.indexPattern()).waitForYellowStatus(),
delegate,
client.admin().cluster()::health
);
})
.<List<InferenceStats>>andThen(
client.threadPool().executor(MachineLearning.UTILITY_THREAD_POOL_NAME),
client.threadPool().getThreadContext(),
(delegate, clusterHealthResponse) -> {
if (clusterHealthResponse.isTimedOut()) {
logger.error(
() -> "[" + Strings.arrayToCommaDelimitedString(modelIds) + "] search failed for models",
response.getFailure()
);
listener.onFailure(
ExceptionsHelper.serverError(
"Searching for stats for models [{}] failed",
response.getFailure(),
Strings.arrayToCommaDelimitedString(modelIds)
)
"getInferenceStats Timed out waiting for index [{}] to be available, "
+ "this will probably cause the request to fail",
MlStatsIndex.indexPattern()
);
return;
}
try {
InferenceStats inferenceStats = handleMultiNodeStatsResponse(response.getResponse(), modelIds[modelIndex++]);
if (inferenceStats != null) {
allStats.add(inferenceStats);
}
} catch (Exception e) {
listener.onFailure(e);

MultiSearchRequest multiSearchRequest = new MultiSearchRequest();
Arrays.stream(modelIds).map(TrainedModelProvider::buildStatsSearchRequest).forEach(multiSearchRequest::add);
if (multiSearchRequest.requests().isEmpty()) {
delegate.onResponse(Collections.emptyList());
return;
}
if (parentTaskId != null) {
multiSearchRequest.setParentTask(parentTaskId);
}
executeAsyncWithOrigin(
client.threadPool().getThreadContext(),
ML_ORIGIN,
multiSearchRequest,
ActionListener.<MultiSearchResponse>wrap(responses -> {
List<InferenceStats> allStats = new ArrayList<>(modelIds.length);
int modelIndex = 0;
assert responses.getResponses().length == modelIds.length
: "mismatch between search response size and models requested";
for (MultiSearchResponse.Item response : responses.getResponses()) {
if (response.isFailure()) {
if (ExceptionsHelper.unwrapCause(response.getFailure()) instanceof ResourceNotFoundException) {
modelIndex++;
continue;
}
logger.error(
() -> "[" + Strings.arrayToCommaDelimitedString(modelIds) + "] search failed for models",
response.getFailure()
);
delegate.onFailure(
ExceptionsHelper.serverError(
"Searching for stats for models [{}] failed",
response.getFailure(),
Strings.arrayToCommaDelimitedString(modelIds)
)
);
return;
}
try {
InferenceStats inferenceStats = handleMultiNodeStatsResponse(
response.getResponse(),
modelIds[modelIndex++]
);
if (inferenceStats != null) {
allStats.add(inferenceStats);
}
} catch (Exception e) {
delegate.onFailure(e);
return;
}
}
delegate.onResponse(allStats);
}, e -> {
Throwable unwrapped = ExceptionsHelper.unwrapCause(e);
if (unwrapped instanceof ResourceNotFoundException) {
delegate.onResponse(Collections.emptyList());
return;
}
delegate.onFailure((Exception) unwrapped);
}),
client::multiSearch
);

}
listener.onResponse(allStats);
}, e -> {
Throwable unwrapped = ExceptionsHelper.unwrapCause(e);
if (unwrapped instanceof ResourceNotFoundException) {
listener.onResponse(Collections.emptyList());
return;
}
listener.onFailure((Exception) unwrapped);
}),
client::multiSearch
);
)
.addListener(
listener,
client.threadPool().executor(MachineLearning.UTILITY_THREAD_POOL_NAME),
client.threadPool().getThreadContext()
);
}

private static SearchRequest buildStatsSearchRequest(String modelId) {
Expand Down
Loading