-
Notifications
You must be signed in to change notification settings - Fork 24.3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Browse files
Browse the repository at this point in the history
* [ML][Inference] PUT API (#50852) This adds the `PUT` API for creating trained models that support our format. This includes * HLRC change for the API * API creation * Validations of model format and call * fixing backport
- Loading branch information
Showing
38 changed files
with
1,638 additions
and
413 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
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
66 changes: 66 additions & 0 deletions
66
client/rest-high-level/src/main/java/org/elasticsearch/client/ml/PutTrainedModelRequest.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,66 @@ | ||
/* | ||
* Licensed to Elasticsearch under one or more contributor | ||
* license agreements. See the NOTICE file distributed with | ||
* this work for additional information regarding copyright | ||
* ownership. Elasticsearch licenses this file to you under | ||
* the Apache License, Version 2.0 (the "License"); you may | ||
* not use this file except in compliance with the License. | ||
* You may obtain a copy of the License at | ||
* | ||
* http://www.apache.org/licenses/LICENSE-2.0 | ||
* | ||
* Unless required by applicable law or agreed to in writing, | ||
* software distributed under the License is distributed on an | ||
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
* KIND, either express or implied. See the License for the | ||
* specific language governing permissions and limitations | ||
* under the License. | ||
*/ | ||
package org.elasticsearch.client.ml; | ||
|
||
import org.elasticsearch.client.Validatable; | ||
import org.elasticsearch.client.ml.inference.TrainedModelConfig; | ||
import org.elasticsearch.common.Strings; | ||
import org.elasticsearch.common.xcontent.ToXContent; | ||
import org.elasticsearch.common.xcontent.ToXContentObject; | ||
import org.elasticsearch.common.xcontent.XContentBuilder; | ||
|
||
import java.io.IOException; | ||
import java.util.Objects; | ||
|
||
|
||
public class PutTrainedModelRequest implements Validatable, ToXContentObject { | ||
|
||
private final TrainedModelConfig config; | ||
|
||
public PutTrainedModelRequest(TrainedModelConfig config) { | ||
this.config = config; | ||
} | ||
|
||
public TrainedModelConfig getTrainedModelConfig() { | ||
return config; | ||
} | ||
|
||
@Override | ||
public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException { | ||
return config.toXContent(builder, params); | ||
} | ||
|
||
@Override | ||
public boolean equals(Object o) { | ||
if (this == o) return true; | ||
if (o == null || getClass() != o.getClass()) return false; | ||
PutTrainedModelRequest request = (PutTrainedModelRequest) o; | ||
return Objects.equals(config, request.config); | ||
} | ||
|
||
@Override | ||
public int hashCode() { | ||
return Objects.hash(config); | ||
} | ||
|
||
@Override | ||
public final String toString() { | ||
return Strings.toString(config); | ||
} | ||
} |
63 changes: 63 additions & 0 deletions
63
...nt/rest-high-level/src/main/java/org/elasticsearch/client/ml/PutTrainedModelResponse.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,63 @@ | ||
/* | ||
* Licensed to Elasticsearch under one or more contributor | ||
* license agreements. See the NOTICE file distributed with | ||
* this work for additional information regarding copyright | ||
* ownership. Elasticsearch licenses this file to you under | ||
* the Apache License, Version 2.0 (the "License"); you may | ||
* not use this file except in compliance with the License. | ||
* You may obtain a copy of the License at | ||
* | ||
* http://www.apache.org/licenses/LICENSE-2.0 | ||
* | ||
* Unless required by applicable law or agreed to in writing, | ||
* software distributed under the License is distributed on an | ||
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
* KIND, either express or implied. See the License for the | ||
* specific language governing permissions and limitations | ||
* under the License. | ||
*/ | ||
package org.elasticsearch.client.ml; | ||
|
||
import org.elasticsearch.client.ml.inference.TrainedModelConfig; | ||
import org.elasticsearch.common.xcontent.ToXContentObject; | ||
import org.elasticsearch.common.xcontent.XContentBuilder; | ||
import org.elasticsearch.common.xcontent.XContentParser; | ||
|
||
import java.io.IOException; | ||
import java.util.Objects; | ||
|
||
|
||
public class PutTrainedModelResponse implements ToXContentObject { | ||
|
||
private final TrainedModelConfig trainedModelConfig; | ||
|
||
public static PutTrainedModelResponse fromXContent(XContentParser parser) throws IOException { | ||
return new PutTrainedModelResponse(TrainedModelConfig.PARSER.parse(parser, null).build()); | ||
} | ||
|
||
public PutTrainedModelResponse(TrainedModelConfig trainedModelConfig) { | ||
this.trainedModelConfig = trainedModelConfig; | ||
} | ||
|
||
public TrainedModelConfig getResponse() { | ||
return trainedModelConfig; | ||
} | ||
|
||
@Override | ||
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { | ||
return trainedModelConfig.toXContent(builder, params); | ||
} | ||
|
||
@Override | ||
public boolean equals(Object o) { | ||
if (this == o) return true; | ||
if (o == null || getClass() != o.getClass()) return false; | ||
PutTrainedModelResponse response = (PutTrainedModelResponse) o; | ||
return Objects.equals(trainedModelConfig, response.trainedModelConfig); | ||
} | ||
|
||
@Override | ||
public int hashCode() { | ||
return Objects.hash(trainedModelConfig); | ||
} | ||
} |
81 changes: 81 additions & 0 deletions
81
...el/src/main/java/org/elasticsearch/client/ml/inference/InferenceToXContentCompressor.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,81 @@ | ||
/* | ||
* Licensed to Elasticsearch under one or more contributor | ||
* license agreements. See the NOTICE file distributed with | ||
* this work for additional information regarding copyright | ||
* ownership. Elasticsearch licenses this file to you under | ||
* the Apache License, Version 2.0 (the "License"); you may | ||
* not use this file except in compliance with the License. | ||
* You may obtain a copy of the License at | ||
* | ||
* http://www.apache.org/licenses/LICENSE-2.0 | ||
* | ||
* Unless required by applicable law or agreed to in writing, | ||
* software distributed under the License is distributed on an | ||
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
* KIND, either express or implied. See the License for the | ||
* specific language governing permissions and limitations | ||
* under the License. | ||
*/ | ||
|
||
package org.elasticsearch.client.ml.inference; | ||
|
||
import org.elasticsearch.common.CheckedFunction; | ||
import org.elasticsearch.common.bytes.BytesArray; | ||
import org.elasticsearch.common.bytes.BytesReference; | ||
import org.elasticsearch.common.io.Streams; | ||
import org.elasticsearch.common.io.stream.BytesStreamOutput; | ||
import org.elasticsearch.common.xcontent.DeprecationHandler; | ||
import org.elasticsearch.common.xcontent.NamedXContentRegistry; | ||
import org.elasticsearch.common.xcontent.ToXContentObject; | ||
import org.elasticsearch.common.xcontent.XContentHelper; | ||
import org.elasticsearch.common.xcontent.XContentParser; | ||
import org.elasticsearch.common.xcontent.XContentType; | ||
|
||
import java.io.IOException; | ||
import java.io.InputStream; | ||
import java.io.OutputStream; | ||
import java.nio.charset.StandardCharsets; | ||
import java.util.Base64; | ||
import java.util.zip.GZIPInputStream; | ||
import java.util.zip.GZIPOutputStream; | ||
|
||
/** | ||
* Collection of helper methods. Similar to CompressedXContent, but this utilizes GZIP. | ||
*/ | ||
public final class InferenceToXContentCompressor { | ||
private static final int BUFFER_SIZE = 4096; | ||
private static final long MAX_INFLATED_BYTES = 1_000_000_000; // 1 gb maximum | ||
|
||
private InferenceToXContentCompressor() {} | ||
|
||
public static <T extends ToXContentObject> String deflate(T objectToCompress) throws IOException { | ||
BytesReference reference = XContentHelper.toXContent(objectToCompress, XContentType.JSON, false); | ||
return deflate(reference); | ||
} | ||
|
||
public static <T> T inflate(String compressedString, | ||
CheckedFunction<XContentParser, T, IOException> parserFunction, | ||
NamedXContentRegistry xContentRegistry) throws IOException { | ||
try(XContentParser parser = XContentHelper.createParser(xContentRegistry, | ||
DeprecationHandler.THROW_UNSUPPORTED_OPERATION, | ||
inflate(compressedString, MAX_INFLATED_BYTES), | ||
XContentType.JSON)) { | ||
return parserFunction.apply(parser); | ||
} | ||
} | ||
|
||
static BytesReference inflate(String compressedString, long streamSize) throws IOException { | ||
byte[] compressedBytes = Base64.getDecoder().decode(compressedString.getBytes(StandardCharsets.UTF_8)); | ||
InputStream gzipStream = new GZIPInputStream(new BytesArray(compressedBytes).streamInput(), BUFFER_SIZE); | ||
InputStream inflateStream = new SimpleBoundedInputStream(gzipStream, streamSize); | ||
return Streams.readFully(inflateStream); | ||
} | ||
|
||
private static String deflate(BytesReference reference) throws IOException { | ||
BytesStreamOutput out = new BytesStreamOutput(); | ||
try (OutputStream compressedOutput = new GZIPOutputStream(out, BUFFER_SIZE)) { | ||
reference.writeTo(compressedOutput); | ||
} | ||
return new String(Base64.getEncoder().encode(BytesReference.toBytes(out.bytes())), StandardCharsets.UTF_8); | ||
} | ||
} |
68 changes: 68 additions & 0 deletions
68
...h-level/src/main/java/org/elasticsearch/client/ml/inference/SimpleBoundedInputStream.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,68 @@ | ||
/* | ||
* Licensed to Elasticsearch under one or more contributor | ||
* license agreements. See the NOTICE file distributed with | ||
* this work for additional information regarding copyright | ||
* ownership. Elasticsearch licenses this file to you under | ||
* the Apache License, Version 2.0 (the "License"); you may | ||
* not use this file except in compliance with the License. | ||
* You may obtain a copy of the License at | ||
* | ||
* http://www.apache.org/licenses/LICENSE-2.0 | ||
* | ||
* Unless required by applicable law or agreed to in writing, | ||
* software distributed under the License is distributed on an | ||
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
* KIND, either express or implied. See the License for the | ||
* specific language governing permissions and limitations | ||
* under the License. | ||
*/ | ||
package org.elasticsearch.client.ml.inference; | ||
|
||
|
||
import java.io.IOException; | ||
import java.io.InputStream; | ||
import java.util.Objects; | ||
|
||
/** | ||
* This is a pared down bounded input stream. | ||
* Only read is specifically enforced. | ||
*/ | ||
final class SimpleBoundedInputStream extends InputStream { | ||
|
||
private final InputStream in; | ||
private final long maxBytes; | ||
private long numBytes; | ||
|
||
SimpleBoundedInputStream(InputStream inputStream, long maxBytes) { | ||
this.in = Objects.requireNonNull(inputStream, "inputStream"); | ||
if (maxBytes < 0) { | ||
throw new IllegalArgumentException("[maxBytes] must be greater than or equal to 0"); | ||
} | ||
this.maxBytes = maxBytes; | ||
} | ||
|
||
|
||
/** | ||
* A simple wrapper around the injected input stream that restricts the total number of bytes able to be read. | ||
* @return The byte read. -1 on internal stream completion or when maxBytes is exceeded. | ||
* @throws IOException on failure | ||
*/ | ||
@Override | ||
public int read() throws IOException { | ||
// We have reached the maximum, signal stream completion. | ||
if (numBytes >= maxBytes) { | ||
return -1; | ||
} | ||
numBytes++; | ||
return in.read(); | ||
} | ||
|
||
/** | ||
* Delegates `close` to the wrapped InputStream | ||
* @throws IOException on failure | ||
*/ | ||
@Override | ||
public void close() throws IOException { | ||
in.close(); | ||
} | ||
} |
Oops, something went wrong.