forked from elastic/elasticsearch
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Adding trained model metadata class. (elastic#106988)
- Loading branch information
Showing
20 changed files
with
948 additions
and
34 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
54 changes: 54 additions & 0 deletions
54
...re/src/main/java/org/elasticsearch/xpack/core/ml/action/FlushTrainedModelCacheAction.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,54 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.core.ml.action; | ||
|
|
||
| import org.elasticsearch.action.ActionType; | ||
| import org.elasticsearch.action.support.master.AcknowledgedRequest; | ||
| import org.elasticsearch.action.support.master.AcknowledgedResponse; | ||
| import org.elasticsearch.common.io.stream.StreamInput; | ||
| import org.elasticsearch.core.TimeValue; | ||
|
|
||
| import java.io.IOException; | ||
| import java.util.Objects; | ||
|
|
||
| public class FlushTrainedModelCacheAction extends ActionType<AcknowledgedResponse> { | ||
|
|
||
| public static final FlushTrainedModelCacheAction INSTANCE = new FlushTrainedModelCacheAction(); | ||
| public static final String NAME = "cluster:admin/xpack/ml/inference/clear_model_cache"; | ||
|
|
||
| private FlushTrainedModelCacheAction() { | ||
| super(NAME); | ||
| } | ||
|
|
||
| public static class Request extends AcknowledgedRequest<FlushTrainedModelCacheAction.Request> { | ||
| public Request() { | ||
| super(); | ||
| } | ||
|
|
||
| Request(TimeValue timeout) { | ||
| super(timeout); | ||
| } | ||
|
|
||
| public Request(StreamInput in) throws IOException { | ||
| super(in); | ||
| } | ||
|
|
||
| @Override | ||
| public int hashCode() { | ||
| return Objects.hashCode(ackTimeout()); | ||
| } | ||
|
|
||
| @Override | ||
| public boolean equals(Object other) { | ||
| if (other == this) return true; | ||
| if (other == null || getClass() != other.getClass()) return false; | ||
| Request that = (Request) other; | ||
| return Objects.equals(that.ackTimeout(), ackTimeout()); | ||
| } | ||
| } | ||
| } |
109 changes: 109 additions & 0 deletions
109
...re/src/main/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelCacheMetadata.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,109 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.core.ml.inference; | ||
|
|
||
| import org.elasticsearch.TransportVersion; | ||
| import org.elasticsearch.TransportVersions; | ||
| import org.elasticsearch.cluster.AbstractNamedDiffable; | ||
| import org.elasticsearch.cluster.ClusterState; | ||
| import org.elasticsearch.cluster.NamedDiff; | ||
| import org.elasticsearch.cluster.metadata.Metadata; | ||
| import org.elasticsearch.common.collect.Iterators; | ||
| import org.elasticsearch.common.io.stream.StreamInput; | ||
| import org.elasticsearch.common.io.stream.StreamOutput; | ||
| import org.elasticsearch.xcontent.ConstructingObjectParser; | ||
| import org.elasticsearch.xcontent.ParseField; | ||
| import org.elasticsearch.xcontent.ToXContent; | ||
| import org.elasticsearch.xcontent.XContentParser; | ||
|
|
||
| import java.io.IOException; | ||
| import java.util.EnumSet; | ||
| import java.util.Iterator; | ||
| import java.util.Objects; | ||
|
|
||
| public class TrainedModelCacheMetadata extends AbstractNamedDiffable<Metadata.Custom> implements Metadata.Custom { | ||
| public static final String NAME = "trained_model_cache_metadata"; | ||
| public static final TrainedModelCacheMetadata EMPTY = new TrainedModelCacheMetadata(0L); | ||
| private static final ParseField VERSION_FIELD = new ParseField("version"); | ||
|
|
||
| @SuppressWarnings("unchecked") | ||
| private static final ConstructingObjectParser<TrainedModelCacheMetadata, Void> PARSER = new ConstructingObjectParser<>( | ||
| NAME, | ||
| true, | ||
| args -> new TrainedModelCacheMetadata((long) args[0]) | ||
| ); | ||
|
|
||
| static { | ||
| PARSER.declareLong(ConstructingObjectParser.constructorArg(), VERSION_FIELD); | ||
| } | ||
|
|
||
| public static TrainedModelCacheMetadata fromXContent(XContentParser parser) { | ||
| return PARSER.apply(parser, null); | ||
| } | ||
|
|
||
| public static TrainedModelCacheMetadata fromState(ClusterState clusterState) { | ||
| TrainedModelCacheMetadata cacheMetadata = clusterState.getMetadata().custom(NAME); | ||
| return cacheMetadata == null ? EMPTY : cacheMetadata; | ||
| } | ||
|
|
||
| public static NamedDiff<Metadata.Custom> readDiffFrom(StreamInput streamInput) throws IOException { | ||
| return readDiffFrom(Metadata.Custom.class, NAME, streamInput); | ||
| } | ||
|
|
||
| private final long version; | ||
|
|
||
| public TrainedModelCacheMetadata(long version) { | ||
| this.version = version; | ||
| } | ||
|
|
||
| public TrainedModelCacheMetadata(StreamInput in) throws IOException { | ||
| this.version = in.readVLong(); | ||
| } | ||
|
|
||
| public long version() { | ||
| return version; | ||
| } | ||
|
|
||
| @Override | ||
| public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params ignored) { | ||
| return Iterators.single(((builder, params) -> { return builder.field(VERSION_FIELD.getPreferredName(), version); })); | ||
| } | ||
|
|
||
| @Override | ||
| public EnumSet<Metadata.XContentContext> context() { | ||
| return Metadata.ALL_CONTEXTS; | ||
| } | ||
|
|
||
| @Override | ||
| public String getWriteableName() { | ||
| return NAME; | ||
| } | ||
|
|
||
| @Override | ||
| public TransportVersion getMinimalSupportedVersion() { | ||
| return TransportVersions.ML_TRAINED_MODEL_CACHE_METADATA_ADDED; | ||
| } | ||
|
|
||
| @Override | ||
| public void writeTo(StreamOutput out) throws IOException { | ||
| out.writeVLong(version); | ||
| } | ||
|
|
||
| @Override | ||
| public boolean equals(Object o) { | ||
| if (this == o) return true; | ||
| if (o == null || getClass() != o.getClass()) return false; | ||
| TrainedModelCacheMetadata that = (TrainedModelCacheMetadata) o; | ||
| return Objects.equals(version, that.version); | ||
| } | ||
|
|
||
| @Override | ||
| public int hashCode() { | ||
| return Objects.hash(version); | ||
| } | ||
| } |
42 changes: 42 additions & 0 deletions
42
...java/org/elasticsearch/xpack/core/ml/action/FlushTrainedModelCacheActionRequestTests.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,42 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.core.ml.action; | ||
|
|
||
| import org.elasticsearch.TransportVersion; | ||
| import org.elasticsearch.common.io.stream.Writeable; | ||
| import org.elasticsearch.core.TimeValue; | ||
| import org.elasticsearch.xpack.core.ml.AbstractBWCWireSerializationTestCase; | ||
| import org.elasticsearch.xpack.core.ml.action.FlushTrainedModelCacheAction.Request; | ||
|
|
||
| import java.io.IOException; | ||
|
|
||
| public class FlushTrainedModelCacheActionRequestTests extends AbstractBWCWireSerializationTestCase<Request> { | ||
| @Override | ||
| protected Request createTestInstance() { | ||
| return randomBoolean() ? new Request() : new Request(randomTimeout()); | ||
| } | ||
|
|
||
| @Override | ||
| protected Request mutateInstance(Request instance) throws IOException { | ||
| return new Request(randomValueOtherThan(instance.timeout(), this::randomTimeout)); | ||
| } | ||
|
|
||
| @Override | ||
| protected Writeable.Reader<Request> instanceReader() { | ||
| return Request::new; | ||
| } | ||
|
|
||
| @Override | ||
| protected Request mutateInstanceForVersion(Request instance, TransportVersion version) { | ||
| return instance; | ||
| } | ||
|
|
||
| private TimeValue randomTimeout() { | ||
| return TimeValue.parseTimeValue(randomTimeValue(), null, "timeout"); | ||
| } | ||
| } |
36 changes: 36 additions & 0 deletions
36
...c/test/java/org/elasticsearch/xpack/core/ml/inference/TrainedModelCacheMetadataTests.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,36 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.core.ml.inference; | ||
|
|
||
| import org.elasticsearch.common.io.stream.Writeable; | ||
| import org.elasticsearch.test.AbstractChunkedSerializingTestCase; | ||
| import org.elasticsearch.xcontent.XContentParser; | ||
|
|
||
| import java.io.IOException; | ||
|
|
||
| public class TrainedModelCacheMetadataTests extends AbstractChunkedSerializingTestCase<TrainedModelCacheMetadata> { | ||
| @Override | ||
| protected TrainedModelCacheMetadata doParseInstance(XContentParser parser) throws IOException { | ||
| return TrainedModelCacheMetadata.fromXContent(parser); | ||
| } | ||
|
|
||
| @Override | ||
| protected Writeable.Reader<TrainedModelCacheMetadata> instanceReader() { | ||
| return TrainedModelCacheMetadata::new; | ||
| } | ||
|
|
||
| @Override | ||
| protected TrainedModelCacheMetadata createTestInstance() { | ||
| return new TrainedModelCacheMetadata(randomNonNegativeLong()); | ||
| } | ||
|
|
||
| @Override | ||
| protected TrainedModelCacheMetadata mutateInstance(TrainedModelCacheMetadata instance) { | ||
| return new TrainedModelCacheMetadata(randomValueOtherThan(instance.version(), () -> randomNonNegativeLong())); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.