Skip to content

Commit a99f1ac

Browse files
valeriy42valeriy42
committed
change function signature
1 parent b1081c0 commit a99f1ac

File tree

10 files changed

+181
-200
lines changed

10 files changed

+181
-200
lines changed

‎x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/TrainedModelAssignmentRebalancer.java‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -135,9 +135,9 @@ private static AssignmentPlan mergePlans(
135135
*/
136136
static void copyAssignments(AssignmentPlan source, AssignmentPlan.Builder dest, Map<String, AssignmentPlan.Node> originalNodeById) {
137137
for (AssignmentPlan.Deployment deployment : source.deployments()) {
138-
Map<AssignmentPlan.Node, Integer> sourceNodeAssignments = source.assignments(deployment).orElse(Map.of());
139-
for (Map.Entry<AssignmentPlan.Node, Integer> sourceAssignment : sourceNodeAssignments.entrySet()) {
140-
AssignmentPlan.Node node = originalNodeById.get(sourceAssignment.getKey().id());
138+
Map<String, Integer> sourceNodeAssignments = source.assignments(deployment).orElse(Map.of());
139+
for (Map.Entry<String, Integer> sourceAssignment : sourceNodeAssignments.entrySet()) {
140+
AssignmentPlan.Node node = originalNodeById.get(sourceAssignment.getKey());
141141
dest.assignModelToNode(deployment, node, sourceAssignment.getValue());
142142
}
143143
}
@@ -337,23 +337,23 @@ private TrainedModelAssignmentMetadata.Builder buildAssignmentsFromPlan(Assignme
337337
assignmentBuilder.setMaxAssignedAllocations(existingAssignment.getMaxAssignedAllocations());
338338
}
339339

340-
Map<AssignmentPlan.Node, Integer> assignments = assignmentPlan.assignments(deployment).orElseGet(Map::of);
341-
for (Map.Entry<AssignmentPlan.Node, Integer> assignment : assignments.entrySet()) {
342-
if (existingAssignment != null && existingAssignment.isRoutedToNode(assignment.getKey().id())) {
343-
RoutingInfo existingRoutingInfo = existingAssignment.getNodeRoutingTable().get(assignment.getKey().id());
340+
Map<String, Integer> assignments = assignmentPlan.assignments(deployment).orElseGet(Map::of);
341+
for (Map.Entry<String, Integer> assignment : assignments.entrySet()) {
342+
if (existingAssignment != null && existingAssignment.isRoutedToNode(assignment.getKey())) {
343+
RoutingInfo existingRoutingInfo = existingAssignment.getNodeRoutingTable().get(assignment.getKey());
344344
RoutingState state = existingRoutingInfo.getState();
345345
String reason = existingRoutingInfo.getReason();
346346
if (state == RoutingState.FAILED) {
347347
state = RoutingState.STARTING;
348348
reason = "";
349349
}
350350
assignmentBuilder.addRoutingEntry(
351-
assignment.getKey().id(),
351+
assignment.getKey(),
352352
new RoutingInfo(existingRoutingInfo.getCurrentAllocations(), assignment.getValue(), state, reason)
353353
);
354354
} else {
355355
assignmentBuilder.addRoutingEntry(
356-
assignment.getKey().id(),
356+
assignment.getKey(),
357357
new RoutingInfo(assignment.getValue(), assignment.getValue(), RoutingState.STARTING, "")
358358
);
359359
}

‎x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AbstractPreserveAllocations.java‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -74,12 +74,9 @@ AssignmentPlan mergePreservedAllocations(AssignmentPlan assignmentPlan) {
7474
// with its preserved allocations.
7575
final Map<Tuple<String, String>, Integer> plannedAssignmentsByDeploymentNodeIdPair = new HashMap<>();
7676
for (Deployment d : assignmentPlan.deployments()) {
77-
Map<Node, Integer> assignmentsOfDeployment = assignmentPlan.assignments(d).orElse(Map.of());
78-
for (Map.Entry<Node, Integer> nodeAssignment : assignmentsOfDeployment.entrySet()) {
79-
plannedAssignmentsByDeploymentNodeIdPair.put(
80-
Tuple.tuple(d.deploymentId(), nodeAssignment.getKey().id()),
81-
nodeAssignment.getValue()
82-
);
77+
Map<String, Integer> assignmentsOfDeployment = assignmentPlan.assignments(d).orElse(Map.of());
78+
for (Map.Entry<String, Integer> nodeAssignment : assignmentsOfDeployment.entrySet()) {
79+
plannedAssignmentsByDeploymentNodeIdPair.put(Tuple.tuple(d.deploymentId(), nodeAssignment.getKey()), nodeAssignment.getValue());
8380
}
8481
}
8582

‎x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AssignmentPlan.java‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -213,9 +213,15 @@ public Set<Deployment> deployments() {
213213
* @param deployment the model for which assignments are returned
214214
* @return the model assignments per node. The Optional will be empty if the model has no assignments.
215215
*/
216-
public Optional<Map<Node, Integer>> assignments(Deployment deployment) {
216+
public Optional<Map<String, Integer>> assignments(Deployment deployment) {
217217
Map<Node, Integer> modelAssignments = assignments.get(deployment);
218-
return (modelAssignments == null || modelAssignments.isEmpty()) ? Optional.empty() : Optional.of(modelAssignments);
218+
if (modelAssignments == null || modelAssignments.isEmpty()) {
219+
return Optional.empty();
220+
}
221+
Map<String, Integer> byNodeId = modelAssignments.entrySet()
222+
.stream()
223+
.collect(Collectors.toMap(e -> e.getKey().id(), Map.Entry::getValue));
224+
return Optional.of(byNodeId);
219225
}
220226

221227
@Override

‎x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AssignmentPlanner.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -157,11 +157,11 @@ private AssignmentPlan solveAllocatingAtLeastOnceModelsThatWerePreviouslyAllocat
157157

158158
Map<String, String> modelIdToNodeIdWithSingleAllocation = new HashMap<>();
159159
for (AssignmentPlan.Deployment m : planWithSingleAllocationForPreviouslyAssignedModels.deployments()) {
160-
Optional<Map<Node, Integer>> assignments = planWithSingleAllocationForPreviouslyAssignedModels.assignments(m);
161-
Set<Node> nodes = assignments.orElse(Map.of()).keySet();
160+
Optional<Map<String, Integer>> assignments = planWithSingleAllocationForPreviouslyAssignedModels.assignments(m);
161+
Set<String> nodes = assignments.orElse(Map.of()).keySet();
162162
if (nodes.isEmpty() == false) {
163163
assert nodes.size() == 1;
164-
modelIdToNodeIdWithSingleAllocation.put(m.deploymentId(), nodes.iterator().next().id());
164+
modelIdToNodeIdWithSingleAllocation.put(m.deploymentId(), nodes.iterator().next());
165165
}
166166
}
167167

‎x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/assignment/planning/ZoneAwareAssignmentPlanner.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -208,9 +208,9 @@ private AssignmentPlan swapOriginalDeploymentsInPlan(
208208

209209
for (AssignmentPlan.Deployment planDeployment : planDeployments) {
210210
AssignmentPlan.Deployment originalDeployment = originalDeploymentsById.get(planDeployment.deploymentId());
211-
Map<Node, Integer> nodeAssignments = plan.assignments(planDeployment).orElse(Map.of());
212-
for (Map.Entry<Node, Integer> assignment : nodeAssignments.entrySet()) {
213-
Node originalNode = originalNodeById.get(assignment.getKey().id());
211+
Map<String, Integer> nodeAssignments = plan.assignments(planDeployment).orElse(Map.of());
212+
for (Map.Entry<String, Integer> assignment : nodeAssignments.entrySet()) {
213+
Node originalNode = originalNodeById.get(assignment.getKey());
214214
finalPlanBuilder.assignModelToNode(originalDeployment, originalNode, assignment.getValue());
215215
}
216216
}
@@ -223,11 +223,11 @@ private Map<String, Map<String, Integer>> mergeAllocationsByNodeIdByDeploymentId
223223
for (AssignmentPlan plan : plans) {
224224
for (AssignmentPlan.Deployment m : plan.deployments()) {
225225
Map<String, Integer> nodeIdToAllocations = allocationsByNodeIdByDeploymentId.get(m.deploymentId());
226-
Optional<Map<Node, Integer>> assignments = plan.assignments(m);
226+
Optional<Map<String, Integer>> assignments = plan.assignments(m);
227227
if (assignments.isPresent()) {
228-
for (Map.Entry<Node, Integer> nodeAssignments : assignments.get().entrySet()) {
228+
for (Map.Entry<String, Integer> nodeAssignments : assignments.get().entrySet()) {
229229
nodeIdToAllocations.compute(
230-
nodeAssignments.getKey().id(),
230+
nodeAssignments.getKey(),
231231
(nodeId, existingAllocations) -> existingAllocations == null
232232
? nodeAssignments.getValue()
233233
: existingAllocations + nodeAssignments.getValue()

‎x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/inference/assignment/planning/AssignmentPlanTests.java‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ public void testAssignModelToNode_GivenNoPreviousAssignment() {
7373

7474
assertThat(plan.deployments(), contains(m));
7575
assertThat(plan.satisfiesCurrentAssignments(), is(true));
76-
assertThat(plan.assignments(m).get(), equalTo(Map.of(n, 1)));
76+
assertThat(plan.assignments(m).get(), equalTo(Map.of(n.id(), 1)));
7777
}
7878
{ // new memory format
7979
AssignmentPlan.Deployment m = new AssignmentPlan.Deployment(
@@ -107,7 +107,7 @@ public void testAssignModelToNode_GivenNoPreviousAssignment() {
107107

108108
assertThat(plan.deployments(), contains(m));
109109
assertThat(plan.satisfiesCurrentAssignments(), is(true));
110-
assertThat(plan.assignments(m).get(), equalTo(Map.of(n, 1)));
110+
assertThat(plan.assignments(m).get(), equalTo(Map.of(n.id(), 1)));
111111
}
112112
}
113113

@@ -140,7 +140,7 @@ public void testAssignModelToNode_GivenNewPlanSatisfiesCurrentAssignment() {
140140

141141
assertThat(plan.deployments(), contains(m));
142142
assertThat(plan.satisfiesCurrentAssignments(), is(true));
143-
assertThat(plan.assignments(m).get(), equalTo(Map.of(n, 1)));
143+
assertThat(plan.assignments(m).get(), equalTo(Map.of(n.id(), 1)));
144144
}
145145
{ // new memory format
146146
AssignmentPlan.Deployment m = new AssignmentPlan.Deployment(
@@ -169,7 +169,7 @@ public void testAssignModelToNode_GivenNewPlanSatisfiesCurrentAssignment() {
169169

170170
assertThat(plan.deployments(), contains(m));
171171
assertThat(plan.satisfiesCurrentAssignments(), is(true));
172-
assertThat(plan.assignments(m).get(), equalTo(Map.of(n, 1)));
172+
assertThat(plan.assignments(m).get(), equalTo(Map.of(n.id(), 1)));
173173

174174
}
175175
}
@@ -195,7 +195,7 @@ public void testAssignModelToNode_GivenNewPlanDoesNotSatisfyCurrentAssignment()
195195

196196
assertThat(plan.deployments(), contains(m));
197197
assertThat(plan.satisfiesCurrentAssignments(), is(false));
198-
assertThat(plan.assignments(m).get(), equalTo(Map.of(n, 1)));
198+
assertThat(plan.assignments(m).get(), equalTo(Map.of(n.id(), 1)));
199199
}
200200
{
201201
// new memory format
@@ -229,7 +229,7 @@ public void testAssignModelToNode_GivenNewPlanDoesNotSatisfyCurrentAssignment()
229229

230230
assertThat(plan.deployments(), contains(m));
231231
assertThat(plan.satisfiesCurrentAssignments(), is(false));
232-
assertThat(plan.assignments(m).get(), equalTo(Map.of(n, 1)));
232+
assertThat(plan.assignments(m).get(), equalTo(Map.of(n.id(), 1)));
233233
}
234234
}
235235

@@ -365,7 +365,7 @@ public void testAssignModelToNode_GivenSameModelAssignedTwice() {
365365

366366
assertThat(plan.deployments(), contains(m));
367367
assertThat(plan.satisfiesCurrentAssignments(), is(true));
368-
assertThat(plan.assignments(m).get(), equalTo(Map.of(n, 3)));
368+
assertThat(plan.assignments(m).get(), equalTo(Map.of(n.id(), 3)));
369369
}
370370

371371
public void testCanAssign_GivenPreviouslyUnassignedModelDoesNotFit() {

0 commit comments

Comments
 (0)