Skip to content

Commit 29d17c0

Browse files
authored
[ML] Fix PyTorchModelIT::testDeploymentStats (#81161)
Adjusts the test's expectations about the information available when deployments are in the `starting` state
1 parent 74797c1 commit 29d17c0

File tree

2 files changed

+19
-17
lines changed

2 files changed

+19
-17
lines changed

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

Lines changed: 19 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -247,18 +247,26 @@ public void testDeploymentStats() throws IOException {
247247
String statusState = (String) XContentMapValues.extractValue("deployment_stats.allocation_status.state", stats.get(0));
248248
assertThat(responseMap.toString(), statusState, is(not(nullValue())));
249249
assertThat(AllocationStatus.State.fromString(statusState), greaterThanOrEqualTo(state));
250-
Integer byteSize = (Integer) XContentMapValues.extractValue("deployment_stats.model_size_bytes", stats.get(0));
251-
assertThat(responseMap.toString(), byteSize, is(not(nullValue())));
252-
assertThat(byteSize, equalTo((int) RAW_MODEL_SIZE));
253250

254-
Response humanResponse = client().performRequest(new Request("GET", "/_ml/trained_models/" + modelId + "/_stats?human"));
255-
var humanResponseMap = entityAsMap(humanResponse);
256-
stats = (List<Map<String, Object>>) humanResponseMap.get("trained_model_stats");
257-
assertThat(stats, hasSize(1));
258-
String stringBytes = (String) XContentMapValues.extractValue("deployment_stats.model_size", stats.get(0));
259-
assertThat("stats response: " + responseMap + " human stats response" + humanResponseMap, stringBytes, is(not(nullValue())));
260-
assertThat(stringBytes, equalTo("1.5kb"));
261-
stopDeployment(model);
251+
// starting models do not know their model size yet
252+
if (state.isAnyOf(AllocationStatus.State.STARTED, AllocationStatus.State.FULLY_ALLOCATED)) {
253+
Integer byteSize = (Integer) XContentMapValues.extractValue("deployment_stats.model_size_bytes", stats.get(0));
254+
assertThat(responseMap.toString(), byteSize, is(not(nullValue())));
255+
assertThat(byteSize, equalTo((int) RAW_MODEL_SIZE));
256+
257+
Response humanResponse = client().performRequest(new Request("GET", "/_ml/trained_models/" + modelId + "/_stats?human"));
258+
var humanResponseMap = entityAsMap(humanResponse);
259+
stats = (List<Map<String, Object>>) humanResponseMap.get("trained_model_stats");
260+
assertThat(stats, hasSize(1));
261+
String stringBytes = (String) XContentMapValues.extractValue("deployment_stats.model_size", stats.get(0));
262+
assertThat(
263+
"stats response: " + responseMap + " human stats response" + humanResponseMap,
264+
stringBytes,
265+
is(not(nullValue()))
266+
);
267+
assertThat(stringBytes, equalTo("1.5kb"));
268+
}
269+
stopDeployment(modelId);
262270
};
263271

264272
assertAtLeast.accept(model, AllocationStatus.State.STARTING);

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

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,6 @@
77

88
package org.elasticsearch.xpack.ml.action;
99

10-
import org.apache.logging.log4j.LogManager;
11-
import org.apache.logging.log4j.Logger;
1210
import org.elasticsearch.action.ActionListener;
1311
import org.elasticsearch.action.FailedNodeException;
1412
import org.elasticsearch.action.TaskOperationFailure;
@@ -56,8 +54,6 @@ public class TransportGetDeploymentStatsAction extends TransportTasksAction<
5654
GetDeploymentStatsAction.Response,
5755
AllocationStats> {
5856

59-
private static final Logger logger = LogManager.getLogger(TransportGetDeploymentStatsAction.class);
60-
6157
@Inject
6258
public TransportGetDeploymentStatsAction(
6359
TransportService transportService,
@@ -274,8 +270,6 @@ static GetDeploymentStatsAction.Response addFailedRoutes(
274270

275271
nodeStats.sort(Comparator.comparing(n -> n.getNode().getId()));
276272

277-
// debug logging added for https://github.com/elastic/elasticsearch/issues/80819
278-
logger.debug("[{}] deployment stats for non-started deployment", modelId);
279273
updatedAllocationStats.add(new AllocationStats(modelId, null, null, null, null, allocation.getStartTime(), nodeStats));
280274
}
281275
}

0 commit comments

Comments
 (0)