AbstractYarnScheduler extends AbstractService implements ResourceScheduler { @Override public N getSchedulerNode(NodeId nodeId) { return nodeTracker.getNode(nodeId); } AbstractYarnScheduler(String name); @Override void serviceInit(Configuration conf); @VisibleForTesting ClusterNodeTracker<N> getNodeTracker(); @VisibleForTesting SchedulingMonitorManager getSchedulingMonitorManager(); List<Container> getTransferredContainers(
      ApplicationAttemptId currentAttempt); Map<ApplicationId, SchedulerApplication<T>> getSchedulerApplications(); List<N> getBlacklistedNodes(final SchedulerApplicationAttempt app); List<N> getNodes(final NodeFilter filter); boolean shouldContainersBeAutoUpdated(); @Override Resource getClusterResource(); @Override Resource getMinimumResourceCapability(); @Override Resource getMaximumResourceCapability(); @Override Resource getMaximumResourceCapability(String queueName); SchedulerHealth getSchedulerHealth(); long getLastNodeUpdateTime(); T getApplicationAttempt(ApplicationAttemptId applicationAttemptId); @Override SchedulerAppReport getSchedulerAppInfo(
      ApplicationAttemptId appAttemptId); @Override ApplicationResourceUsageReport getAppResourceUsageReport(
      ApplicationAttemptId appAttemptId); T getCurrentAttemptForContainer(ContainerId containerId); @Override RMContainer getRMContainer(ContainerId containerId); @Override SchedulerNodeReport getNodeReport(NodeId nodeId); @Override String moveApplication(ApplicationId appId, String newQueue); @Override void preValidateMoveApplication(ApplicationId appId,
      String newQueue); void removeQueue(String queueName); @Override void addQueue(Queue newQueue); @Override void setEntitlement(String queue, QueueEntitlement entitlement); void recoverContainersOnNode(List<NMContainerStatus> containerReports,
      RMNode nm); @VisibleForTesting void clearPendingContainerCache(); @VisibleForTesting @Private // clean up a completed container void completedContainer(RMContainer rmContainer,
      ContainerStatus containerStatus, RMContainerEventType event); @Override N getSchedulerNode(NodeId nodeId); @Override void moveAllApps(String sourceQueue, String destQueue); @Override void killAllAppsInQueue(String queueName); void updateNodeResource(RMNode nm,
      ResourceOption resourceOption); @Override EnumSet<SchedulerResourceTypes> getSchedulingResourceTypes(); @Override Set<String> getPlanQueues(); boolean placementConstraintEnabled(); @Override List<ResourceRequest> getPendingResourceRequestsForAttempt(
      ApplicationAttemptId attemptId); @Override List<SchedulingRequest> getPendingSchedulingRequestsForAttempt(
      ApplicationAttemptId attemptId); @Override Priority checkAndGetApplicationPriority(
          Priority priorityRequestedByApp, UserGroupInformation user,
          String queuePath, ApplicationId applicationId); @Override Priority updateApplicationPriority(Priority newPriority,
      ApplicationId applicationId, SettableFuture<Object> future,
      UserGroupInformation user); @Override Priority getMaxClusterLevelAppPriority(); @Override void setClusterMaxPriority(Configuration conf); ActivitiesManager getActivitiesManager(); Clock getClock(); @VisibleForTesting void setClock(Clock clock); @Lock(Lock.NoLock.class) SchedulerNode getNode(NodeId nodeId); @Override Resource getNormalizedResource(Resource requestedResource,
                                        Resource maxResourceCapability); @Override List<NodeId> getNodeIds(String resourceName); void asyncContainerRelease(RMContainer container); Resource getMinimumAllocation(); Resource getMaximumAllocation(); @Override long checkAndGetApplicationLifetime(String queueName, long lifetime); @Override long getMaximumApplicationLifetime(String queueName); @VisibleForTesting abstract void killContainer(RMContainer container); @VisibleForTesting void update(); @Override void reinitialize(Configuration conf, RMContext rmContext); @Override boolean attemptAllocationOnNode(SchedulerApplicationAttempt appAttempt,
      SchedulingRequest schedulingRequest, SchedulerNode schedulerNode); @Override void resetSchedulerMetrics();  }
@Test(timeout = 30000l) public void testContainerReleaseWithAllocationTags() throws Exception { if (getSchedulerType().equals(SchedulerType.CAPACITY)) { final String testTag1 = "some-tag"; final String testTag2 = "some-other-tag"; YarnConfiguration conf = getConf(); conf.set(YarnConfiguration.RM_PLACEMENT_CONSTRAINTS_HANDLER, "scheduler"); MockRM rm1 = new MockRM(conf); rm1.start(); MockNM nm1 = new MockNM("127.0.0.1:1234", 10240, rm1.getResourceTrackerService()); nm1.registerNode(); MockRMAppSubmissionData data = MockRMAppSubmissionData.Builder.createWithMemory(200, rm1) .withAppName("name") .withUser("user") .withAcls(new HashMap<>()) .withUnmanagedAM(false) .withQueue("default") .withMaxAppAttempts(-1) .withCredentials(null) .withAppType("Test") .withWaitForAppAcceptedState(false) .withKeepContainers(true) .build(); RMApp app1 = MockRMAppSubmitter.submit(rm1, data); MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1); SchedulingRequest sr = SchedulingRequest .newInstance(1l, Priority.newInstance(1), ExecutionTypeRequest.newInstance(ExecutionType.GUARANTEED), Sets.newHashSet(testTag1), ResourceSizing.newInstance(1, Resource.newInstance(1024, 1)), null); SchedulingRequest sr1 = SchedulingRequest .newInstance(2l, Priority.newInstance(1), ExecutionTypeRequest.newInstance(ExecutionType.GUARANTEED), Sets.newHashSet(testTag2), ResourceSizing.newInstance(3, Resource.newInstance(1024, 1)), null); AllocateRequest ar = AllocateRequest.newBuilder() .schedulingRequests(Lists.newArrayList(sr, sr1)).build(); am1.allocate(ar); nm1.nodeHeartbeat(true); List<Container> allocated = new ArrayList<>(); while (allocated.size() < 4) { AllocateResponse rsp = am1 .allocate(new ArrayList<>(), new ArrayList<>()); allocated.addAll(rsp.getAllocatedContainers()); nm1.nodeHeartbeat(true); Thread.sleep(1000); } Assert.assertEquals(4, allocated.size()); Set<Container> containers = allocated.stream() .filter(container -> container.getAllocationRequestId() == 1l) .collect(Collectors.toSet()); Assert.assertNotNull(containers); Assert.assertEquals(1, containers.size()); ContainerId cid = containers.iterator().next().getId(); rm1.getRMContext().getScheduler() .getSchedulerNode(nm1.getNodeId()).containerStarted(cid); Map<String, Long> nodeTags = rm1.getRMContext() .getAllocationTagsManager() .getAllocationTagsWithCount(nm1.getNodeId()); Assert.assertNotNull(nodeTags.get(testTag1)); Assert.assertEquals(1, nodeTags.get(testTag1).intValue()); am1.allocate(new ArrayList<>(), Lists.newArrayList(cid)); nodeTags = rm1.getRMContext().getAllocationTagsManager() .getAllocationTagsWithCount(nm1.getNodeId()); Assert.assertNotNull(nodeTags); Assert.assertNotNull(nodeTags.get(testTag1)); Assert.assertEquals(1, nodeTags.get(testTag1).intValue()); ContainerStatus cs = ContainerStatus.newInstance(cid, ContainerState.COMPLETE, "", 0); nm1.nodeHeartbeat(Lists.newArrayList(cs), true); GenericTestUtils.waitFor(() -> { Map<String, Long> tags = rm1.getRMContext() .getAllocationTagsManager() .getAllocationTagsWithCount(nm1.getNodeId()); return tags.get(testTag1) == null && tags.get(testTag2).intValue() == 3; }, 500, 3000); } }
@Test(timeout = 30000L) public void testNodeRemovedWithAllocationTags() throws Exception { if (getSchedulerType().equals(SchedulerType.CAPACITY)) { final String testTag1 = "some-tag"; YarnConfiguration conf = getConf(); conf.set(YarnConfiguration.RM_PLACEMENT_CONSTRAINTS_HANDLER, "scheduler"); MockRM rm1 = new MockRM(conf); rm1.start(); MockNM nm1 = new MockNM("127.0.0.1:1234", 10240, rm1.getResourceTrackerService()); nm1.registerNode(); MockRMAppSubmissionData data = MockRMAppSubmissionData.Builder.createWithMemory(200, rm1) .withAppName("name") .withUser("user") .withAcls(new HashMap<>()) .withUnmanagedAM(false) .withQueue("default") .withMaxAppAttempts(-1) .withCredentials(null) .withAppType("Test") .withWaitForAppAcceptedState(false) .withKeepContainers(true) .build(); RMApp app1 = MockRMAppSubmitter.submit(rm1, data); MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1); SchedulingRequest sr = SchedulingRequest .newInstance(1L, Priority.newInstance(1), ExecutionTypeRequest.newInstance(ExecutionType.GUARANTEED), Sets.newHashSet(testTag1), ResourceSizing.newInstance(1, Resource.newInstance(1024, 1)), null); AllocateRequest ar = AllocateRequest.newBuilder() .schedulingRequests(Lists.newArrayList(sr)).build(); am1.allocate(ar); nm1.nodeHeartbeat(true); List<Container> allocated = new ArrayList<>(); while (allocated.size() < 1) { AllocateResponse rsp = am1 .allocate(new ArrayList<>(), new ArrayList<>()); allocated.addAll(rsp.getAllocatedContainers()); nm1.nodeHeartbeat(true); Thread.sleep(1000); } Assert.assertEquals(1, allocated.size()); Set<Container> containers = allocated.stream() .filter(container -> container.getAllocationRequestId() == 1L) .collect(Collectors.toSet()); Assert.assertNotNull(containers); Assert.assertEquals(1, containers.size()); ContainerId cid = containers.iterator().next().getId(); rm1.getRMContext().getScheduler() .getSchedulerNode(nm1.getNodeId()).containerStarted(cid); Map<String, Long> nodeTags = rm1.getRMContext() .getAllocationTagsManager() .getAllocationTagsWithCount(nm1.getNodeId()); Assert.assertNotNull(nodeTags.get(testTag1)); Assert.assertEquals(1, nodeTags.get(testTag1).intValue()); RMNode node1 = MockNodes.newNodeInfo( 0, Resources.createResource(nm1.getMemory()), 1, "127.0.0.1", 1234); rm1.getRMContext().getScheduler().handle( new NodeRemovedSchedulerEvent(node1)); nodeTags = rm1.getRMContext().getAllocationTagsManager() .getAllocationTagsWithCount(nm1.getNodeId()); Assert.assertNull(nodeTags); } }