Browse Source

[ML] Fix PyTorchModelIT::testDeploymentStats (#81161)

Adjusts the test's expectations about the information available
when deployments are in the `starting` state
David Kyle 3 years ago
parent
commit
29d17c0460

+ 19 - 11
x-pack/plugin/ml/qa/native-multi-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/PyTorchModelIT.java

@@ -247,18 +247,26 @@ public class PyTorchModelIT extends ESRestTestCase {
             String statusState = (String) XContentMapValues.extractValue("deployment_stats.allocation_status.state", stats.get(0));
             assertThat(responseMap.toString(), statusState, is(not(nullValue())));
             assertThat(AllocationStatus.State.fromString(statusState), greaterThanOrEqualTo(state));
-            Integer byteSize = (Integer) XContentMapValues.extractValue("deployment_stats.model_size_bytes", stats.get(0));
-            assertThat(responseMap.toString(), byteSize, is(not(nullValue())));
-            assertThat(byteSize, equalTo((int) RAW_MODEL_SIZE));
 
-            Response humanResponse = client().performRequest(new Request("GET", "/_ml/trained_models/" + modelId + "/_stats?human"));
-            var humanResponseMap = entityAsMap(humanResponse);
-            stats = (List<Map<String, Object>>) humanResponseMap.get("trained_model_stats");
-            assertThat(stats, hasSize(1));
-            String stringBytes = (String) XContentMapValues.extractValue("deployment_stats.model_size", stats.get(0));
-            assertThat("stats response: " + responseMap + " human stats response" + humanResponseMap, stringBytes, is(not(nullValue())));
-            assertThat(stringBytes, equalTo("1.5kb"));
-            stopDeployment(model);
+            // starting models do not know their model size yet
+            if (state.isAnyOf(AllocationStatus.State.STARTED, AllocationStatus.State.FULLY_ALLOCATED)) {
+                Integer byteSize = (Integer) XContentMapValues.extractValue("deployment_stats.model_size_bytes", stats.get(0));
+                assertThat(responseMap.toString(), byteSize, is(not(nullValue())));
+                assertThat(byteSize, equalTo((int) RAW_MODEL_SIZE));
+
+                Response humanResponse = client().performRequest(new Request("GET", "/_ml/trained_models/" + modelId + "/_stats?human"));
+                var humanResponseMap = entityAsMap(humanResponse);
+                stats = (List<Map<String, Object>>) humanResponseMap.get("trained_model_stats");
+                assertThat(stats, hasSize(1));
+                String stringBytes = (String) XContentMapValues.extractValue("deployment_stats.model_size", stats.get(0));
+                assertThat(
+                    "stats response: " + responseMap + " human stats response" + humanResponseMap,
+                    stringBytes,
+                    is(not(nullValue()))
+                );
+                assertThat(stringBytes, equalTo("1.5kb"));
+            }
+            stopDeployment(modelId);
         };
 
         assertAtLeast.accept(model, AllocationStatus.State.STARTING);

+ 0 - 6
x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportGetDeploymentStatsAction.java

@@ -7,8 +7,6 @@
 
 package org.elasticsearch.xpack.ml.action;
 
-import org.apache.logging.log4j.LogManager;
-import org.apache.logging.log4j.Logger;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.FailedNodeException;
 import org.elasticsearch.action.TaskOperationFailure;
@@ -56,8 +54,6 @@ public class TransportGetDeploymentStatsAction extends TransportTasksAction<
     GetDeploymentStatsAction.Response,
     AllocationStats> {
 
-    private static final Logger logger = LogManager.getLogger(TransportGetDeploymentStatsAction.class);
-
     @Inject
     public TransportGetDeploymentStatsAction(
         TransportService transportService,
@@ -274,8 +270,6 @@ public class TransportGetDeploymentStatsAction extends TransportTasksAction<
 
                 nodeStats.sort(Comparator.comparing(n -> n.getNode().getId()));
 
-                // debug logging added for https://github.com/elastic/elasticsearch/issues/80819
-                logger.debug("[{}] deployment stats for non-started deployment", modelId);
                 updatedAllocationStats.add(new AllocationStats(modelId, null, null, null, null, allocation.getStartTime(), nodeStats));
             }
         }