Some review comments

elastic · Apr 2, 2021 · b3cc616 · b3cc616
1 parent 4badd50
commit b3cc616
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 6 deletions.
diff --git a/...c/main/java/org/elasticsearch/xpack/core/ml/action/InferTrainedModelDeploymentAction.java b/...c/main/java/org/elasticsearch/xpack/core/ml/action/InferTrainedModelDeploymentAction.java
@@ -129,5 +129,11 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             result.toXContent(builder, params);
             return builder;
         }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            super.writeTo(out);
+            result.writeTo(out);
+        }
     }
 }
diff --git a/...ore/src/main/java/org/elasticsearch/xpack/core/ml/inference/deployment/PyTorchResult.java b/...ore/src/main/java/org/elasticsearch/xpack/core/ml/inference/deployment/PyTorchResult.java
@@ -46,11 +46,7 @@ public class PyTorchResult implements ToXContentObject, Writeable {
                 double[][] primitiveDoubles = new double[listOfListOfDoubles.size()][];
                 for (int i = 0; i < listOfListOfDoubles.size(); i++) {
                     List<Double> row = listOfListOfDoubles.get(i);
-                    double[] primitiveRow = new double[row.size()];
-                    for (int j = 0; j < row.size(); j++) {
-                        primitiveRow[j] = row.get(j);
-                    }
-                    primitiveDoubles[i] = primitiveRow;
+                    primitiveDoubles[i] = row.stream().mapToDouble(d -> d).toArray();
                 }
                 return primitiveDoubles;
             },

diff --git a/...n/ml/src/main/java/org/elasticsearch/xpack/ml/inference/deployment/DeploymentManager.java b/...n/ml/src/main/java/org/elasticsearch/xpack/ml/inference/deployment/DeploymentManager.java
@@ -11,11 +11,13 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.message.ParameterizedMessage;
 import org.apache.lucene.util.SetOnce;
+import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.get.GetRequest;
 import org.elasticsearch.action.get.GetResponse;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.common.unit.TimeValue;
+import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.core.ml.inference.deployment.PyTorchResult;
 import org.elasticsearch.xpack.core.ml.inference.deployment.TrainedModelDeploymentState;
@@ -113,7 +115,8 @@ private void waitForResult(ProcessContext processContext, String requestId, Acti
             TimeValue timeout = TimeValue.timeValueSeconds(5);
             PyTorchResult pyTorchResult = processContext.resultProcessor.waitForResult(requestId, timeout);
             if (pyTorchResult == null) {
-                listener.onFailure(ExceptionsHelper.serverError("no result was produced within timeout value [{}]", timeout));
+                listener.onFailure(new ElasticsearchStatusException("timeout [{}] waiting for inference result",
+                    RestStatus.TOO_MANY_REQUESTS, timeout));
             } else {
                 listener.onResponse(pyTorchResult);
             }