tidy up

elastic · Oct 14, 2024 · e649ffb · e649ffb
1 parent 26e32bb
commit e649ffb
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 22 deletions.
diff --git a/...tTest/java/org/elasticsearch/xpack/ml/integration/AdaptiveAllocationsScaleFromZeroIT.java b/...tTest/java/org/elasticsearch/xpack/ml/integration/AdaptiveAllocationsScaleFromZeroIT.java
@@ -96,15 +96,6 @@ public void onFailure(Exception exception) {
         if (failures.isEmpty() == false) {
             fail(failures.getFirst());
         }
-
-        // {
-        // var responseMap = entityAsMap(getTrainedModelStats(modelId));
-        // List<Map<String, Object>> stats = (List<Map<String, Object>>) responseMap.get("trained_model_stats");
-        // String statusState = (String) XContentMapValues.extractValue("deployment_stats.allocation_status.state", stats.get(0));
-        // assertThat(responseMap.toString(), statusState, is(not(nullValue())));
-        // Integer count = (Integer) XContentMapValues.extractValue("deployment_stats.allocation_status.allocation_count", stats.get(0));
-        // assertThat(responseMap.toString(), count, greaterThan(0));
-        // }
     }
 
     @SuppressWarnings("unchecked")
@@ -136,7 +127,7 @@ public void testMultipleDeploymentsWaiting() throws Exception {
         }, 30, TimeUnit.SECONDS);
 
         // infer will scale up
-        int inferenceCount = 100;
+        int inferenceCount = 10;
         var latch = new CountDownLatch(inferenceCount);
         for (int i = 0; i < inferenceCount; i++) {
             asyncInfer("Auto scale and infer", randomFrom(idsList), TimeValue.timeValueSeconds(5), new ResponseListener() {

diff --git a/...rc/javaRestTest/java/org/elasticsearch/xpack/ml/integration/PyTorchModelRestTestCase.java b/...rc/javaRestTest/java/org/elasticsearch/xpack/ml/integration/PyTorchModelRestTestCase.java
@@ -12,11 +12,8 @@
 import org.elasticsearch.client.Response;
 import org.elasticsearch.client.ResponseListener;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
-import org.elasticsearch.common.xcontent.ChunkedToXContentHelper;
-import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
 import org.elasticsearch.common.xcontent.support.XContentMapValues;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.test.SecuritySettingsSourceField;
@@ -299,12 +296,6 @@ protected Response startDeployment(String modelId, String deploymentId, Adaptive
             + "&threads_per_allocation=1"
             + "&wait_for=started";
 
-        ChunkedToXContentObject innerChunkedContent = params -> Iterators.concat(
-            ChunkedToXContentHelper.startObject(),
-            Iterators.single(((builder, p2) -> builder.field("adaptive_allocations", adaptiveAllocationsSettings))),
-            ChunkedToXContentHelper.endObject()
-        );
-
         XContentBuilder builder = JsonXContent.contentBuilder();
         builder.startObject();
         builder.field("adaptive_allocations", adaptiveAllocationsSettings);

diff --git a/...ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java b/...ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportInternalInferModelAction.java
@@ -293,6 +293,7 @@ private void inferAgainstAllocatedModel(
 
         assert nodes.stream().mapToInt(Tuple::v2).sum() == request.numberOfDocuments()
             : "mismatch; sum of node requests does not match number of documents in request";
+        inferOnAssignmentNodes(nodes, request, responseBuilder, parentTaskId, listener);
     }
 
     private void inferOnBlockedRequest(InferenceWaitForAllocation.WaitingRequest request, TrainedModelAssignment assignment) {

diff --git a/...gin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/InferenceWaitForAllocation.java b/...gin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/InferenceWaitForAllocation.java
@@ -97,13 +97,13 @@ public synchronized void waitForAssignment(WaitingRequest request) {
         }
 
         pendingRequestCount.incrementAndGet();
-        var prediate = new DeploymentHasAtLeastOneAllocation(request.deploymentId());
+        var predicate = new DeploymentHasAtLeastOneAllocation(request.deploymentId());
 
         assignmentService.waitForAssignmentCondition(
             request.deploymentId(),
-            prediate,
+            predicate,
             request.request().getInferenceTimeout(),
-            new WaitingListener(request.deploymentId(), request, prediate)
+            new WaitingListener(request.deploymentId(), request, predicate)
         );
     }