[ML][Inference] PUT API

client/rest-high-level/src/main/java/org/elasticsearch/client/MLRequestConverters.java

@@ -73,6 +73,7 @@
 import org.elasticsearch.client.ml.PutDatafeedRequest;
 import org.elasticsearch.client.ml.PutFilterRequest;
 import org.elasticsearch.client.ml.PutJobRequest;
+import org.elasticsearch.client.ml.PutTrainedModelRequest;
 import org.elasticsearch.client.ml.RevertModelSnapshotRequest;
 import org.elasticsearch.client.ml.SetUpgradeModeRequest;
 import org.elasticsearch.client.ml.StartDataFrameAnalyticsRequest;
@@ -792,6 +793,16 @@ static Request deleteTrainedModel(DeleteTrainedModelRequest deleteRequest) {
         return new Request(HttpDelete.METHOD_NAME, endpoint);
     }
 
+    static Request putTrainedModel(PutTrainedModelRequest putTrainedModelRequest) throws IOException {
+        String endpoint = new EndpointBuilder()
+            .addPathPartAsIs("_ml", "inference")
+            .addPathPart(putTrainedModelRequest.getTrainedModelConfig().getModelId())
+            .build();
+        Request request = new Request(HttpPut.METHOD_NAME, endpoint);
+        request.setEntity(createEntity(putTrainedModelRequest, REQUEST_BODY_CONTENT_TYPE));
+        return request;
+    }
+
     static Request putFilter(PutFilterRequest putFilterRequest) throws IOException {
         String endpoint = new EndpointBuilder()
             .addPathPartAsIs("_ml")

client/rest-high-level/src/main/java/org/elasticsearch/client/MachineLearningClient.java

@@ -100,6 +100,8 @@
 import org.elasticsearch.client.ml.PutFilterResponse;
 import org.elasticsearch.client.ml.PutJobRequest;
 import org.elasticsearch.client.ml.PutJobResponse;
+import org.elasticsearch.client.ml.PutTrainedModelRequest;
+import org.elasticsearch.client.ml.PutTrainedModelResponse;
 import org.elasticsearch.client.ml.RevertModelSnapshotRequest;
 import org.elasticsearch.client.ml.RevertModelSnapshotResponse;
 import org.elasticsearch.client.ml.SetUpgradeModeRequest;
@@ -2340,6 +2342,48 @@ public Cancellable getTrainedModelsAsync(GetTrainedModelsRequest request,
             Collections.emptySet());
     }
 
+    /**
+     * Put trained model config
+     * <p>
+     * For additional info
+     * see <a href="TODO">
+     *     PUT Trained Model Config documentation</a>
+     *
+     * @param request The {@link PutTrainedModelRequest}
+     * @param options Additional request options (e.g. headers), use {@link RequestOptions#DEFAULT} if nothing needs to be customized
+     * @return {@link PutTrainedModelResponse} response object
+     */
+    public PutTrainedModelResponse putTrainedModel(PutTrainedModelRequest request, RequestOptions options) throws IOException {
+        return restHighLevelClient.performRequestAndParseEntity(request,
+            MLRequestConverters::putTrainedModel,
+            options,
+            PutTrainedModelResponse::fromXContent,
+            Collections.emptySet());
+    }
+
+    /**
+     * Put trained model config asynchronously and notifies listener upon completion
+     * <p>
+     * For additional info
+     * see <a href="TODO">
+     *     PUT Trained Model Config documentation</a>
+     *
+     * @param request The {@link PutTrainedModelRequest}
+     * @param options Additional request options (e.g. headers), use {@link RequestOptions#DEFAULT} if nothing needs to be customized
+     * @param listener Listener to be notified upon request completion
+     * @return cancellable that may be used to cancel the request
+     */
+    public Cancellable putTrainedModelAsync(PutTrainedModelRequest request,
+                                            RequestOptions options,
+                                            ActionListener<PutTrainedModelResponse> listener) {
+        return restHighLevelClient.performRequestAsyncAndParseEntity(request,
+            MLRequestConverters::putTrainedModel,
+            options,
+            PutTrainedModelResponse::fromXContent,
+            listener,
+            Collections.emptySet());
+    }
+
     /**
      * Gets trained model stats
      * <p>

client/rest-high-level/src/main/java/org/elasticsearch/client/ml/PutTrainedModelRequest.java

@@ -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);
+    }
+}

client/rest-high-level/src/main/java/org/elasticsearch/client/ml/PutTrainedModelResponse.java

@@ -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);
+    }
+}

client/rest-high-level/src/main/java/org/elasticsearch/client/ml/inference/InferenceToXContentCompressor.java

@@ -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);
+    }
+}

client/rest-high-level/src/main/java/org/elasticsearch/client/ml/inference/SimpleBoundedInputStream.java

@@ -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();
+    }
+}