diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/allocator/RegularContainerAllocator.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/allocator/RegularContainerAllocator.java index f53ece51ffe4b3..fb4452b37cb364 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/allocator/RegularContainerAllocator.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/allocator/RegularContainerAllocator.java @@ -848,6 +848,7 @@ private ContainerAllocation allocate(Resource clusterResource, // Directly return if this check fails. ContainerAllocation result; ContainerAllocation lastReservation = null; + FiCaSchedulerNode lastReservationNode = null; AppPlacementAllocator schedulingPS = application.getAppSchedulingInfo().getAppPlacementAllocator( @@ -911,8 +912,7 @@ private ContainerAllocation allocate(Resource clusterResource, resourceLimits, schedulerKey, reservedContainer); if (AllocationState.ALLOCATED == result.getAllocationState()) { - result = doAllocation(result, node, schedulerKey, reservedContainer); - break; + return doAllocation(result, node, schedulerKey, reservedContainer); } // In MultiNodePlacement, Try Allocate on other Available nodes @@ -921,15 +921,21 @@ private ContainerAllocation allocate(Resource clusterResource, // iterator could not fit and returns RESERVED allocation. if (AllocationState.RESERVED == result.getAllocationState()) { lastReservation = result; + lastReservationNode = node; if (iter.hasNext()) { continue; } else { - result = doAllocation(lastReservation, node, schedulerKey, - reservedContainer); + return doAllocation(lastReservation, lastReservationNode, + schedulerKey, reservedContainer); } } } + if (lastReservation != null) { + return doAllocation(lastReservation, lastReservationNode, schedulerKey, + reservedContainer); + } + return result; } diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/allocator/TestRegularContainerAllocator.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/allocator/TestRegularContainerAllocator.java new file mode 100644 index 00000000000000..20be316b6caf9f --- /dev/null +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/allocator/TestRegularContainerAllocator.java @@ -0,0 +1,230 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.allocator; + +import java.lang.reflect.Method; +import java.util.Arrays; +import java.util.Iterator; + +import org.apache.hadoop.yarn.api.records.NodeId; +import org.apache.hadoop.yarn.api.records.Resource; +import org.apache.hadoop.yarn.conf.YarnConfiguration; +import org.apache.hadoop.yarn.server.resourcemanager.RMContext; +import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainer; +import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AppSchedulingInfo; +import org.apache.hadoop.yarn.server.resourcemanager.scheduler.ResourceLimits; +import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.SchedulingMode; +import org.apache.hadoop.yarn.server.resourcemanager.scheduler.common.fica.FiCaSchedulerApp; +import org.apache.hadoop.yarn.server.resourcemanager.scheduler.common.fica.FiCaSchedulerNode; +import org.apache.hadoop.yarn.server.resourcemanager.scheduler.placement.AppPlacementAllocator; +import org.apache.hadoop.yarn.server.resourcemanager.scheduler.placement.CandidateNodeSet; +import org.apache.hadoop.yarn.server.scheduler.SchedulerRequestKey; +import org.apache.hadoop.yarn.util.resource.DefaultResourceCalculator; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class TestRegularContainerAllocator { + + @Test + public void testReservationFallbackUsesReservationNode() throws Exception { + FiCaSchedulerApp application = mock(FiCaSchedulerApp.class); + RMContext rmContext = mock(RMContext.class); + when(rmContext.getYarnConfiguration()).thenReturn(new YarnConfiguration()); + + SchedulerRequestKey schedulerKey = mock(SchedulerRequestKey.class); + CandidateNodeSet candidates = mock(CandidateNodeSet.class); + ResourceLimits resourceLimits = mock(ResourceLimits.class); + RMContainer reservedContainer = mock(RMContainer.class); + + FiCaSchedulerNode reservationNode = mock(FiCaSchedulerNode.class); + FiCaSchedulerNode skippedNode = mock(FiCaSchedulerNode.class); + when(reservationNode.getNodeID()) + .thenReturn(NodeId.newInstance("reservation", 1234)); + when(skippedNode.getNodeID()) + .thenReturn(NodeId.newInstance("skipped", 1234)); + + mockPlacement(application, schedulerKey, candidates, + reservationNode, skippedNode); + when(application.getOutstandingAsksCount(schedulerKey)).thenReturn(1); + + ContainerAllocation reservation = new ContainerAllocation(null, + Resource.newInstance(1024, 1), AllocationState.RESERVED); + TestAllocator allocator = new TestAllocator(application, rmContext, + reservation, ContainerAllocation.APP_SKIPPED); + + ContainerAllocation result = invokeAllocate(allocator, candidates, + resourceLimits, schedulerKey, reservedContainer); + + assertEquals(AllocationState.RESERVED, result.getAllocationState()); + assertSame(reservationNode, allocator.getLastDoAllocationNode()); + } + + @Test + public void testLaterReservationReplacesEarlierReservation() + throws Exception { + FiCaSchedulerApp application = mock(FiCaSchedulerApp.class); + RMContext rmContext = mock(RMContext.class); + when(rmContext.getYarnConfiguration()).thenReturn(new YarnConfiguration()); + + SchedulerRequestKey schedulerKey = mock(SchedulerRequestKey.class); + CandidateNodeSet candidates = mock(CandidateNodeSet.class); + ResourceLimits resourceLimits = mock(ResourceLimits.class); + RMContainer reservedContainer = mock(RMContainer.class); + + FiCaSchedulerNode firstReservationNode = mock(FiCaSchedulerNode.class); + FiCaSchedulerNode secondReservationNode = mock(FiCaSchedulerNode.class); + FiCaSchedulerNode skippedNode = mock(FiCaSchedulerNode.class); + when(firstReservationNode.getNodeID()) + .thenReturn(NodeId.newInstance("reservation1", 1234)); + when(secondReservationNode.getNodeID()) + .thenReturn(NodeId.newInstance("reservation2", 1234)); + when(skippedNode.getNodeID()) + .thenReturn(NodeId.newInstance("skipped", 1234)); + + mockPlacement(application, schedulerKey, candidates, + firstReservationNode, secondReservationNode, skippedNode); + when(application.getOutstandingAsksCount(schedulerKey)).thenReturn(1); + + ContainerAllocation firstReservation = new ContainerAllocation(null, + Resource.newInstance(1024, 1), AllocationState.RESERVED); + ContainerAllocation secondReservation = new ContainerAllocation(null, + Resource.newInstance(2048, 1), AllocationState.RESERVED); + TestAllocator allocator = new TestAllocator(application, rmContext, + firstReservation, secondReservation, ContainerAllocation.APP_SKIPPED); + + ContainerAllocation result = invokeAllocate(allocator, candidates, + resourceLimits, schedulerKey, reservedContainer); + + assertSame(secondReservation, result); + assertSame(secondReservationNode, allocator.getLastDoAllocationNode()); + } + + @Test + public void testFailedAllocationPreventsReservationFallback() + throws Exception { + FiCaSchedulerApp application = mock(FiCaSchedulerApp.class); + RMContext rmContext = mock(RMContext.class); + when(rmContext.getYarnConfiguration()).thenReturn(new YarnConfiguration()); + + SchedulerRequestKey schedulerKey = mock(SchedulerRequestKey.class); + CandidateNodeSet candidates = mock(CandidateNodeSet.class); + ResourceLimits resourceLimits = mock(ResourceLimits.class); + RMContainer reservedContainer = mock(RMContainer.class); + + FiCaSchedulerNode reservationNode = mock(FiCaSchedulerNode.class); + FiCaSchedulerNode allocationNode = mock(FiCaSchedulerNode.class); + when(reservationNode.getNodeID()) + .thenReturn(NodeId.newInstance("reservation", 1234)); + when(allocationNode.getNodeID()) + .thenReturn(NodeId.newInstance("allocation", 1234)); + + mockPlacement(application, schedulerKey, candidates, + reservationNode, allocationNode); + when(application.getOutstandingAsksCount(schedulerKey)).thenReturn(1); + + ContainerAllocation reservation = new ContainerAllocation(null, + Resource.newInstance(1024, 1), AllocationState.RESERVED); + ContainerAllocation allocation = new ContainerAllocation(null, + Resource.newInstance(1024, 1), AllocationState.ALLOCATED); + TestAllocator allocator = new TestAllocator(application, rmContext, + reservation, allocation); + allocator.setDoAllocationResult(ContainerAllocation.APP_SKIPPED); + + ContainerAllocation result = invokeAllocate(allocator, candidates, + resourceLimits, schedulerKey, reservedContainer); + + assertEquals(AllocationState.APP_SKIPPED, result.getAllocationState()); + assertSame(allocationNode, allocator.getLastDoAllocationNode()); + } + + private void mockPlacement(FiCaSchedulerApp application, + SchedulerRequestKey schedulerKey, + CandidateNodeSet candidates, + FiCaSchedulerNode... nodes) { + AppPlacementAllocator placement = + mock(AppPlacementAllocator.class); + AppSchedulingInfo appSchedulingInfo = mock(AppSchedulingInfo.class); + when(application.getAppSchedulingInfo()).thenReturn(appSchedulingInfo); + when(appSchedulingInfo.getAppPlacementAllocator( + schedulerKey)).thenReturn(placement); + when(placement.getPreferredNodeIterator(candidates)) + .thenReturn(Arrays.asList(nodes).iterator()); + } + + private ContainerAllocation invokeAllocate( + RegularContainerAllocator allocator, + CandidateNodeSet candidates, + ResourceLimits resourceLimits, + SchedulerRequestKey schedulerKey, + RMContainer reservedContainer) throws Exception { + Method allocate = RegularContainerAllocator.class.getDeclaredMethod( + "allocate", Resource.class, CandidateNodeSet.class, + SchedulingMode.class, ResourceLimits.class, SchedulerRequestKey.class, + RMContainer.class); + allocate.setAccessible(true); + return (ContainerAllocation) allocate.invoke(allocator, + Resource.newInstance(8192, 8), candidates, + SchedulingMode.RESPECT_PARTITION_EXCLUSIVITY, resourceLimits, + schedulerKey, reservedContainer); + } + + private static final class TestAllocator extends RegularContainerAllocator { + private final Iterator results; + private ContainerAllocation doAllocationResult; + private FiCaSchedulerNode lastDoAllocationNode; + + private TestAllocator(FiCaSchedulerApp application, RMContext rmContext, + ContainerAllocation... results) { + super(application, new DefaultResourceCalculator(), rmContext, null); + this.results = Arrays.asList(results).iterator(); + } + + private void setDoAllocationResult(ContainerAllocation result) { + this.doAllocationResult = result; + } + + private FiCaSchedulerNode getLastDoAllocationNode() { + return lastDoAllocationNode; + } + + @Override + ContainerAllocation tryAllocateOnNode(Resource clusterResource, + FiCaSchedulerNode node, SchedulingMode schedulingMode, + ResourceLimits resourceLimits, SchedulerRequestKey schedulerKey, + RMContainer reservedContainer) { + return results.next(); + } + + @Override + ContainerAllocation doAllocation(ContainerAllocation allocationResult, + FiCaSchedulerNode node, SchedulerRequestKey schedulerKey, + RMContainer reservedContainer) { + lastDoAllocationNode = node; + if (doAllocationResult != null) { + return doAllocationResult; + } + allocationResult.updatedContainer = reservedContainer; + return allocationResult; + } + } +}