Skip to content
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

Jump to bottom

[ML] make more GET and heavier processing APIs cancellable #88142

Merged
merged 6 commits into from
Jul 5, 2022
Merged

[ML] make more GET and heavier processing APIs cancellable #88142

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
bd28f76
[ML] make more GET and heavier processing APIs cancellable
benwtrent Jun 28, 2022
ae8f0b7
fixing tests, adding bwc test for explain
benwtrent Jun 29, 2022
742ab70
Merge branch 'master' into feature/ml-make-heavy-apis-cancellable
elasticmachine Jun 29, 2022
e7fea43
Merge branch 'master' into feature/ml-make-heavy-apis-cancellable
elasticmachine Jul 5, 2022
fd8400a
Apply suggestions from code review
benwtrent Jul 5, 2022
59959cb
Merge branch 'master' into feature/ml-make-heavy-apis-cancellable
elasticmachine Jul 5, 2022
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 ...in/core/src/main/java/org/elasticsearch/xpack/core/ml/action/EvaluateDataFrameAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,9 @@
import org.elasticsearch.common.xcontent.XContentParserUtils;
import org.elasticsearch.core.Nullable;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.tasks.CancellableTask;
import org.elasticsearch.tasks.Task;
import org.elasticsearch.tasks.TaskId;
import org.elasticsearch.xcontent.ConstructingObjectParser;
import org.elasticsearch.xcontent.ParseField;
import org.elasticsearch.xcontent.ToXContentObject;
Expand All @@ -30,6 +33,7 @@
import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;

Expand Down Expand Up @@ -175,6 +179,11 @@ public boolean equals(Object o) {
&& Objects.equals(queryProvider, that.queryProvider)
&& Objects.equals(evaluation, that.evaluation);
}

@Override
public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
return new CancellableTask(id, type, action, "evaluate_data_frame", parentTaskId, headers);
}
}

public static class Response extends ActionResponse implements ToXContentObject {
Expand Down
127 changes: 127 additions & 0 deletions ...src/main/java/org/elasticsearch/xpack/core/ml/action/ExplainDataFrameAnalyticsAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,21 +6,36 @@
*/
package org.elasticsearch.xpack.core.ml.action;

import org.elasticsearch.action.ActionRequestValidationException;
import org.elasticsearch.action.ActionResponse;
import org.elasticsearch.action.ActionType;
import org.elasticsearch.action.ValidateActions;
import org.elasticsearch.action.support.master.AcknowledgedRequest;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.tasks.CancellableTask;
import org.elasticsearch.tasks.Task;
import org.elasticsearch.tasks.TaskId;
import org.elasticsearch.xcontent.ConstructingObjectParser;
import org.elasticsearch.xcontent.ParseField;
import org.elasticsearch.xcontent.ToXContentObject;
import org.elasticsearch.xcontent.XContentBuilder;
import org.elasticsearch.xcontent.XContentParser;
import org.elasticsearch.xpack.core.common.validation.SourceDestValidator;
import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsConfig;
import org.elasticsearch.xpack.core.ml.dataframe.DataFrameAnalyticsSource;
import org.elasticsearch.xpack.core.ml.dataframe.explain.FieldSelection;
import org.elasticsearch.xpack.core.ml.dataframe.explain.MemoryEstimation;
import org.elasticsearch.xpack.core.ml.job.messages.Messages;
import org.elasticsearch.xpack.core.ml.utils.MlStrings;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.Objects;

import static org.elasticsearch.core.Strings.format;

public class ExplainDataFrameAnalyticsAction extends ActionType<ExplainDataFrameAnalyticsAction.Response> {

public static final ExplainDataFrameAnalyticsAction INSTANCE = new ExplainDataFrameAnalyticsAction();
Expand All @@ -30,6 +45,118 @@ private ExplainDataFrameAnalyticsAction() {
super(NAME, ExplainDataFrameAnalyticsAction.Response::new);
}

public static class Request extends AcknowledgedRequest<Request> implements ToXContentObject {
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 added this request object (just a copy of PutDataFrameAnalytics) as the PUT object was being reused here, but I didn't want to make that request's task cancellable (as PUT isn't cancellable).

Copy link
Member Author

Choose a reason for hiding this comment

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

Also, I added a BWC test for explain in a mixed cluster to ensure this doesn't break anything. Ran it many many times and its all good since the explain request and the put request can be serialized to each other without issue.

public static Request parseRequest(XContentParser parser) {
DataFrameAnalyticsConfig.Builder configBuilder = DataFrameAnalyticsConfig.STRICT_PARSER.apply(parser, null);
DataFrameAnalyticsConfig config = configBuilder.buildForExplain();
return new Request(config);
}

private final DataFrameAnalyticsConfig config;

public Request(StreamInput in) throws IOException {
super(in);
config = new DataFrameAnalyticsConfig(in);
}

public Request(DataFrameAnalyticsConfig config) {
this.config = config;
}

@Override
public void writeTo(StreamOutput out) throws IOException {
super.writeTo(out);
config.writeTo(out);
}

public DataFrameAnalyticsConfig getConfig() {
return config;
}

@Override
public ActionRequestValidationException validate() {
ActionRequestValidationException error = null;
error = checkConfigIdIsValid(config, error);
error = SourceDestValidator.validateRequest(error, config.getDest().getIndex());
error = checkNoIncludedAnalyzedFieldsAreExcludedBySourceFiltering(config, error);
return error;
}

private ActionRequestValidationException checkConfigIdIsValid(
DataFrameAnalyticsConfig analyticsConfig,
ActionRequestValidationException error
) {
if (MlStrings.isValidId(analyticsConfig.getId()) == false) {
error = ValidateActions.addValidationError(
Messages.getMessage(Messages.INVALID_ID, DataFrameAnalyticsConfig.ID, analyticsConfig.getId()),
error
);
}
if (MlStrings.hasValidLengthForId(analyticsConfig.getId()) == false) {
error = ValidateActions.addValidationError(
Messages.getMessage(
Messages.ID_TOO_LONG,
DataFrameAnalyticsConfig.ID,
analyticsConfig.getId(),
MlStrings.ID_LENGTH_LIMIT
),
error
);
}
return error;
}

private ActionRequestValidationException checkNoIncludedAnalyzedFieldsAreExcludedBySourceFiltering(
DataFrameAnalyticsConfig analyticsConfig,
ActionRequestValidationException error
) {
if (analyticsConfig.getAnalyzedFields() == null) {
return error;
}
for (String analyzedInclude : analyticsConfig.getAnalyzedFields().includes()) {
if (analyticsConfig.getSource().isFieldExcluded(analyzedInclude)) {
return ValidateActions.addValidationError(
"field ["
+ analyzedInclude
+ "] is included in ["
+ DataFrameAnalyticsConfig.ANALYZED_FIELDS.getPreferredName()
+ "] but not in ["
+ DataFrameAnalyticsConfig.SOURCE.getPreferredName()
+ "."
+ DataFrameAnalyticsSource._SOURCE.getPreferredName()
+ "]",
error
);
}
}
return error;
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
config.toXContent(builder, params);
return builder;
}

@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
Request request = (Request) o;
return Objects.equals(config, request.config);
}

@Override
public int hashCode() {
return Objects.hash(config);
}

@Override
public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
return new CancellableTask(id, type, action, format("explain_data_frame_analytics[%s]", config.getId()), parentTaskId, headers);
}
}

public static class Response extends ActionResponse implements ToXContentObject {

public static final ParseField TYPE = new ParseField("explain_data_frame_analytics_response");
Expand Down
10 changes: 10 additions & 0 deletions ...plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetCategoriesAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,9 @@
import org.elasticsearch.action.ActionType;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.tasks.CancellableTask;
import org.elasticsearch.tasks.Task;
import org.elasticsearch.tasks.TaskId;
import org.elasticsearch.xcontent.ObjectParser;
import org.elasticsearch.xcontent.ParseField;
import org.elasticsearch.xcontent.ToXContentObject;
Expand All @@ -24,9 +27,11 @@
import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;

import java.io.IOException;
import java.util.Map;
import java.util.Objects;

import static org.elasticsearch.action.ValidateActions.addValidationError;
import static org.elasticsearch.core.Strings.format;

public class GetCategoriesAction extends ActionType<GetCategoriesAction.Response> {

Expand Down Expand Up @@ -190,6 +195,11 @@ public boolean equals(Object o) {
public int hashCode() {
return Objects.hash(jobId, categoryId, pageParams, partitionFieldValue);
}

@Override
public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
return new CancellableTask(id, type, action, format("get_categories[%s:%s]", jobId, categoryId), parentTaskId, headers);
}
}

public static class Response extends AbstractGetResourcesResponse<CategoryDefinition> implements ToXContentObject {
Expand Down
11 changes: 11 additions & 0 deletions .../plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetDatafeedsAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,21 @@
import org.elasticsearch.action.support.master.MasterNodeReadRequest;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.tasks.CancellableTask;
import org.elasticsearch.tasks.Task;
import org.elasticsearch.tasks.TaskId;
import org.elasticsearch.xcontent.ToXContentObject;
import org.elasticsearch.xpack.core.action.AbstractGetResourcesResponse;
import org.elasticsearch.xpack.core.action.util.QueryPage;
import org.elasticsearch.xpack.core.ml.datafeed.DatafeedConfig;
import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;

import java.io.IOException;
import java.util.Map;
import java.util.Objects;

import static org.elasticsearch.core.Strings.format;

public class GetDatafeedsAction extends ActionType<GetDatafeedsAction.Response> {

public static final GetDatafeedsAction INSTANCE = new GetDatafeedsAction();
Expand Down Expand Up @@ -93,6 +99,11 @@ public boolean equals(Object obj) {
Request other = (Request) obj;
return Objects.equals(datafeedId, other.datafeedId) && Objects.equals(allowNoMatch, other.allowNoMatch);
}

@Override
public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
return new CancellableTask(id, type, action, format("get_datafeeds[%s]", datafeedId), parentTaskId, headers);
}
}

public static class Response extends AbstractGetResourcesResponse<DatafeedConfig> implements ToXContentObject {
Expand Down
11 changes: 11 additions & 0 deletions x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetJobsAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,21 @@
import org.elasticsearch.action.support.master.MasterNodeReadRequest;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.tasks.CancellableTask;
import org.elasticsearch.tasks.Task;
import org.elasticsearch.tasks.TaskId;
import org.elasticsearch.xcontent.ToXContentObject;
import org.elasticsearch.xpack.core.action.AbstractGetResourcesResponse;
import org.elasticsearch.xpack.core.action.util.QueryPage;
import org.elasticsearch.xpack.core.ml.job.config.Job;
import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;

import java.io.IOException;
import java.util.Map;
import java.util.Objects;

import static org.elasticsearch.core.Strings.format;

public class GetJobsAction extends ActionType<GetJobsAction.Response> {

public static final GetJobsAction INSTANCE = new GetJobsAction();
Expand Down Expand Up @@ -91,6 +97,11 @@ public boolean equals(Object obj) {
Request other = (Request) obj;
return Objects.equals(jobId, other.jobId) && Objects.equals(allowNoMatch, other.allowNoMatch);
}

@Override
public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
return new CancellableTask(id, type, action, format("get_anomaly_detection_jobs[%s]", jobId), parentTaskId, headers);
}
}

public static class Response extends AbstractGetResourcesResponse<Job> implements ToXContentObject {
Expand Down
11 changes: 11 additions & 0 deletions ...in/core/src/main/java/org/elasticsearch/xpack/core/ml/action/GetModelSnapshotsAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.core.Nullable;
import org.elasticsearch.tasks.CancellableTask;
import org.elasticsearch.tasks.Task;
import org.elasticsearch.tasks.TaskId;
import org.elasticsearch.xcontent.ObjectParser;
import org.elasticsearch.xcontent.ParseField;
import org.elasticsearch.xcontent.ToXContentObject;
Expand All @@ -25,8 +28,11 @@
import org.elasticsearch.xpack.core.ml.utils.ExceptionsHelper;

import java.io.IOException;
import java.util.Map;
import java.util.Objects;

import static org.elasticsearch.core.Strings.format;

public class GetModelSnapshotsAction extends ActionType<GetModelSnapshotsAction.Response> {

public static final GetModelSnapshotsAction INSTANCE = new GetModelSnapshotsAction();
Expand Down Expand Up @@ -205,6 +211,11 @@ public boolean equals(Object obj) {
&& Objects.equals(sort, other.sort)
&& Objects.equals(desc, other.desc);
}

@Override
public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
return new CancellableTask(id, type, action, format("get_job_model_snapshot[%s:%s]", jobId, snapshotId), parentTaskId, headers);
}
}

public static class Response extends AbstractGetResourcesResponse<ModelSnapshot> implements ToXContentObject {
Expand Down
10 changes: 10 additions & 0 deletions ...src/main/java/org/elasticsearch/xpack/core/ml/action/PreviewDataFrameAnalyticsAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@
import org.elasticsearch.action.ActionType;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.tasks.CancellableTask;
import org.elasticsearch.tasks.Task;
import org.elasticsearch.tasks.TaskId;
import org.elasticsearch.xcontent.ConstructingObjectParser;
import org.elasticsearch.xcontent.ObjectParser;
import org.elasticsearch.xcontent.ParseField;
Expand All @@ -26,6 +29,8 @@
import java.util.Map;
import java.util.Objects;

import static org.elasticsearch.core.Strings.format;

public class PreviewDataFrameAnalyticsAction extends ActionType<PreviewDataFrameAnalyticsAction.Response> {

public static final PreviewDataFrameAnalyticsAction INSTANCE = new PreviewDataFrameAnalyticsAction();
Expand Down Expand Up @@ -87,6 +92,11 @@ public int hashCode() {
return Objects.hash(config);
}

@Override
public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
return new CancellableTask(id, type, action, format("preview_data_frame_analytics[%s]", config.getId()), parentTaskId, headers);
}

public static class Builder {
private DataFrameAnalyticsConfig config;

Expand Down
10 changes: 10 additions & 0 deletions ...ugin/core/src/main/java/org/elasticsearch/xpack/core/ml/action/PreviewDatafeedAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,9 @@
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.core.Nullable;
import org.elasticsearch.tasks.CancellableTask;
import org.elasticsearch.tasks.Task;
import org.elasticsearch.tasks.TaskId;
import org.elasticsearch.xcontent.ObjectParser;
import org.elasticsearch.xcontent.ParseField;
import org.elasticsearch.xcontent.ToXContentObject;
Expand All @@ -29,9 +32,11 @@

import java.io.IOException;
import java.io.InputStream;
import java.util.Map;
import java.util.Objects;
import java.util.OptionalLong;

import static org.elasticsearch.core.Strings.format;
import static org.elasticsearch.xpack.core.ml.action.StartDatafeedAction.DatafeedParams.parseDateOrThrow;
import static org.elasticsearch.xpack.core.ml.action.StartDatafeedAction.END_TIME;
import static org.elasticsearch.xpack.core.ml.action.StartDatafeedAction.START_TIME;
Expand Down Expand Up @@ -199,6 +204,11 @@ public boolean equals(Object obj) {
&& Objects.equals(jobConfig, other.jobConfig);
}

@Override
public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
return new CancellableTask(id, type, action, format("preview_datafeed[%s]", datafeedId), parentTaskId, headers);
}

public static class Builder {
private String datafeedId;
private DatafeedConfig.Builder datafeedBuilder;
Expand Down
Loading