diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/assignment/AssignmentStats.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/assignment/AssignmentStats.java index 31b513eea161e..a518654ec6785 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/assignment/AssignmentStats.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/assignment/AssignmentStats.java @@ -432,6 +432,22 @@ public int hashCode() { private final Instant startTime; private final List nodeStats; + public AssignmentStats(AssignmentStats other) { + this.deploymentId = other.deploymentId; + this.modelId = other.modelId; + this.threadsPerAllocation = other.threadsPerAllocation; + this.numberOfAllocations = other.numberOfAllocations; + this.adaptiveAllocationsSettings = other.adaptiveAllocationsSettings; + this.queueCapacity = other.queueCapacity; + this.startTime = other.startTime; + this.nodeStats = other.nodeStats; + this.state = other.state; + this.reason = other.reason; + this.allocationStatus = other.allocationStatus; + this.cacheSize = other.cacheSize; + this.priority = other.priority; + } + public AssignmentStats( String deploymentId, String modelId, @@ -535,6 +551,12 @@ public AssignmentState getState() { return state; } + public AssignmentStats setNodeStats(List nodeStats) { + this.nodeStats.clear(); + this.nodeStats.addAll(nodeStats); + return this; + } + public AssignmentStats setState(AssignmentState state) { this.state = state; return this; diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/assignment/AssignmentStatsTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/assignment/AssignmentStatsTests.java index 07c56b073cd00..3b327b480e075 100644 --- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/assignment/AssignmentStatsTests.java +++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/inference/assignment/AssignmentStatsTests.java @@ -214,6 +214,12 @@ public void testGetOverallInferenceStatsWithOnlyStoppedNodes() { assertThat(stats.getFailureCount(), equalTo(0L)); } + public void testCopyConstructor() { + AssignmentStats original = randomDeploymentStats(); + AssignmentStats copy = new AssignmentStats(original); + assertThat(copy, equalTo(original)); + } + @Override protected Writeable.Reader instanceReader() { return AssignmentStats::new;