List of usage examples for org.apache.hadoop.yarn.api.records NodeId getHost
@Public @Stable public abstract String getHost();
From source file:org.apache.tez.dag.app.rm.node.TestAMNodeTracker.java
License:Apache License
private static NodeReport generateNodeReport(NodeId nodeId, NodeState nodeState) { NodeReport nodeReport = mock(NodeReport.class); doReturn(nodeId).when(nodeReport).getNodeId(); doReturn(nodeState).when(nodeReport).getNodeState(); String httpAddress = nodeId.getHost() + ":3433"; doReturn(httpAddress).when(nodeReport).getHttpAddress(); doReturn("/default-rack").when(nodeReport).getRackName(); doReturn(Resource.newInstance(0, 0)).when(nodeReport).getUsed(); doReturn(Resource.newInstance(10240, 12)).when(nodeReport).getCapability(); doReturn(10).when(nodeReport).getNumContainers(); doReturn(nodeState.toString()).when(nodeReport).getHealthReport(); long healthReportTime = System.currentTimeMillis(); doReturn(healthReportTime).when(nodeReport).getLastHealthReportTime(); return nodeReport; }
From source file:org.apache.tez.dag.app.rm.TestTaskScheduler.java
License:Apache License
@SuppressWarnings({ "unchecked" }) @Test(timeout = 10000)/*from w w w . j ava2 s. c o m*/ public void testTaskSchedulerNoReuse() throws Exception { RackResolver.init(new YarnConfiguration()); TaskSchedulerAppCallback mockApp = mock(TaskSchedulerAppCallback.class); AppContext mockAppContext = mock(AppContext.class); when(mockAppContext.getAMState()).thenReturn(DAGAppMasterState.RUNNING); TezAMRMClientAsync<CookieContainerRequest> mockRMClient = mock(TezAMRMClientAsync.class); String appHost = "host"; int appPort = 0; String appUrl = "url"; TaskSchedulerWithDrainableAppCallback scheduler = new TaskSchedulerWithDrainableAppCallback(mockApp, new AlwaysMatchesContainerMatcher(), appHost, appPort, appUrl, mockRMClient, mockAppContext); TaskSchedulerAppCallbackDrainable drainableAppCallback = scheduler.getDrainableAppCallback(); Configuration conf = new Configuration(); conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false); int interval = 100; conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, interval); scheduler.init(conf); drainableAppCallback.drain(); verify(mockRMClient).init(conf); verify(mockRMClient).setHeartbeatInterval(interval); RegisterApplicationMasterResponse mockRegResponse = mock(RegisterApplicationMasterResponse.class); Resource mockMaxResource = mock(Resource.class); Map<ApplicationAccessType, String> mockAcls = mock(Map.class); when(mockRegResponse.getMaximumResourceCapability()).thenReturn(mockMaxResource); when(mockRegResponse.getApplicationACLs()).thenReturn(mockAcls); ByteBuffer mockKey = mock(ByteBuffer.class); when(mockRegResponse.getClientToAMTokenMasterKey()).thenReturn(mockKey); when(mockRMClient.registerApplicationMaster(anyString(), anyInt(), anyString())) .thenReturn(mockRegResponse); scheduler.start(); drainableAppCallback.drain(); verify(mockRMClient).start(); verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl); verify(mockApp).setApplicationRegistrationData(mockMaxResource, mockAcls, mockKey); when(mockRMClient.getClusterNodeCount()).thenReturn(5); Assert.assertEquals(5, scheduler.getClusterNodeCount()); Resource mockClusterResource = mock(Resource.class); when(mockRMClient.getAvailableResources()).thenReturn(mockClusterResource); Assert.assertEquals(mockClusterResource, mockRMClient.getAvailableResources()); Object mockTask1 = mock(Object.class); Object mockCookie1 = mock(Object.class); Resource mockCapability = mock(Resource.class); String[] hosts = { "host1", "host5" }; String[] racks = { "/default-rack", "/default-rack" }; Priority mockPriority = mock(Priority.class); ArgumentCaptor<CookieContainerRequest> requestCaptor = ArgumentCaptor .forClass(CookieContainerRequest.class); // allocate task scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1); drainableAppCallback.drain(); verify(mockRMClient, times(1)).addContainerRequest((CookieContainerRequest) any()); // returned from task requests before allocation happens assertFalse(scheduler.deallocateTask(mockTask1, true)); verify(mockApp, times(0)).containerBeingReleased(any(ContainerId.class)); verify(mockRMClient, times(1)).removeContainerRequest((CookieContainerRequest) any()); verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any()); // deallocating unknown task assertFalse(scheduler.deallocateTask(mockTask1, true)); verify(mockApp, times(0)).containerBeingReleased(any(ContainerId.class)); verify(mockRMClient, times(1)).removeContainerRequest((CookieContainerRequest) any()); verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any()); // allocate tasks Object mockTask2 = mock(Object.class); Object mockCookie2 = mock(Object.class); Object mockTask3 = mock(Object.class); Object mockCookie3 = mock(Object.class); scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1); drainableAppCallback.drain(); verify(mockRMClient, times(2)).addContainerRequest(requestCaptor.capture()); CookieContainerRequest request1 = requestCaptor.getValue(); scheduler.allocateTask(mockTask2, mockCapability, hosts, racks, mockPriority, null, mockCookie2); drainableAppCallback.drain(); verify(mockRMClient, times(3)).addContainerRequest(requestCaptor.capture()); CookieContainerRequest request2 = requestCaptor.getValue(); scheduler.allocateTask(mockTask3, mockCapability, hosts, racks, mockPriority, null, mockCookie3); drainableAppCallback.drain(); verify(mockRMClient, times(4)).addContainerRequest(requestCaptor.capture()); CookieContainerRequest request3 = requestCaptor.getValue(); List<Container> containers = new ArrayList<Container>(); Container mockContainer1 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer1.getNodeId().getHost()).thenReturn("host1"); ContainerId mockCId1 = mock(ContainerId.class); when(mockContainer1.getId()).thenReturn(mockCId1); containers.add(mockContainer1); Container mockContainer2 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer2.getNodeId().getHost()).thenReturn("host2"); ContainerId mockCId2 = mock(ContainerId.class); when(mockContainer2.getId()).thenReturn(mockCId2); containers.add(mockContainer2); Container mockContainer3 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer3.getNodeId().getHost()).thenReturn("host3"); ContainerId mockCId3 = mock(ContainerId.class); when(mockContainer3.getId()).thenReturn(mockCId3); containers.add(mockContainer3); Container mockContainer4 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer4.getNodeId().getHost()).thenReturn("host4"); ContainerId mockCId4 = mock(ContainerId.class); when(mockContainer4.getId()).thenReturn(mockCId4); containers.add(mockContainer4); ArrayList<CookieContainerRequest> hostContainers = new ArrayList<CookieContainerRequest>(); hostContainers.add(request1); hostContainers.add(request2); hostContainers.add(request3); ArrayList<CookieContainerRequest> rackContainers = new ArrayList<CookieContainerRequest>(); rackContainers.add(request2); rackContainers.add(request3); ArrayList<CookieContainerRequest> anyContainers = new ArrayList<CookieContainerRequest>(); anyContainers.add(request3); final List<ArrayList<CookieContainerRequest>> hostList = new LinkedList<ArrayList<CookieContainerRequest>>(); hostList.add(hostContainers); final List<ArrayList<CookieContainerRequest>> rackList = new LinkedList<ArrayList<CookieContainerRequest>>(); rackList.add(rackContainers); final List<ArrayList<CookieContainerRequest>> anyList = new LinkedList<ArrayList<CookieContainerRequest>>(); anyList.add(anyContainers); final List<ArrayList<CookieContainerRequest>> emptyList = new LinkedList<ArrayList<CookieContainerRequest>>(); // return all requests for host1 when(mockRMClient.getMatchingRequests((Priority) any(), eq("host1"), (Resource) any())) .thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return hostList; } }); // first request matched by host // second request matched to rack. RackResolver by default puts hosts in // /default-rack. We need to workaround by returning rack matches only once when(mockRMClient.getMatchingRequests((Priority) any(), eq("/default-rack"), (Resource) any())) .thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return rackList; } }).thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return emptyList; } }); // third request matched to ANY when(mockRMClient.getMatchingRequests((Priority) any(), eq(ResourceRequest.ANY), (Resource) any())) .thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return anyList; } }).thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return emptyList; } }); scheduler.onContainersAllocated(containers); drainableAppCallback.drain(); // first container allocated verify(mockApp).taskAllocated(mockTask1, mockCookie1, mockContainer1); verify(mockApp).taskAllocated(mockTask2, mockCookie2, mockContainer2); verify(mockApp).taskAllocated(mockTask3, mockCookie3, mockContainer3); // no other allocations returned verify(mockApp, times(3)).taskAllocated(any(), any(), (Container) any()); verify(mockRMClient).removeContainerRequest(request1); verify(mockRMClient).removeContainerRequest(request2); verify(mockRMClient).removeContainerRequest(request3); // verify unwanted container released verify(mockRMClient).releaseAssignedContainer(mockCId4); // deallocate allocated task assertTrue(scheduler.deallocateTask(mockTask1, true)); drainableAppCallback.drain(); verify(mockApp).containerBeingReleased(mockCId1); verify(mockRMClient).releaseAssignedContainer(mockCId1); // deallocate allocated container Assert.assertEquals(mockTask2, scheduler.deallocateContainer(mockCId2)); drainableAppCallback.drain(); verify(mockRMClient).releaseAssignedContainer(mockCId2); verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any()); List<ContainerStatus> statuses = new ArrayList<ContainerStatus>(); ContainerStatus mockStatus1 = mock(ContainerStatus.class); when(mockStatus1.getContainerId()).thenReturn(mockCId1); statuses.add(mockStatus1); ContainerStatus mockStatus2 = mock(ContainerStatus.class); when(mockStatus2.getContainerId()).thenReturn(mockCId2); statuses.add(mockStatus2); ContainerStatus mockStatus3 = mock(ContainerStatus.class); when(mockStatus3.getContainerId()).thenReturn(mockCId3); statuses.add(mockStatus3); ContainerStatus mockStatus4 = mock(ContainerStatus.class); when(mockStatus4.getContainerId()).thenReturn(mockCId4); statuses.add(mockStatus4); scheduler.onContainersCompleted(statuses); drainableAppCallback.drain(); // released container status returned verify(mockApp).containerCompleted(mockTask1, mockStatus1); verify(mockApp).containerCompleted(mockTask2, mockStatus2); // currently allocated container status returned and not released verify(mockApp).containerCompleted(mockTask3, mockStatus3); // no other statuses returned verify(mockApp, times(3)).containerCompleted(any(), (ContainerStatus) any()); verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any()); // verify blacklisting verify(mockRMClient, times(0)).addNodeToBlacklist((NodeId) any()); String badHost = "host6"; NodeId badNodeId = mock(NodeId.class); when(badNodeId.getHost()).thenReturn(badHost); scheduler.blacklistNode(badNodeId); verify(mockRMClient, times(1)).addNodeToBlacklist(badNodeId); Object mockTask4 = mock(Object.class); Object mockCookie4 = mock(Object.class); scheduler.allocateTask(mockTask4, mockCapability, null, null, mockPriority, null, mockCookie4); drainableAppCallback.drain(); verify(mockRMClient, times(5)).addContainerRequest(requestCaptor.capture()); CookieContainerRequest request4 = requestCaptor.getValue(); anyContainers.clear(); anyContainers.add(request4); Container mockContainer5 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer5.getNodeId().getHost()).thenReturn(badHost); when(mockContainer5.getNodeId()).thenReturn(badNodeId); ContainerId mockCId5 = mock(ContainerId.class); when(mockContainer5.getId()).thenReturn(mockCId5); containers.clear(); containers.add(mockContainer5); when(mockRMClient.getMatchingRequests((Priority) any(), eq(ResourceRequest.ANY), (Resource) any())) .thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return anyList; } }).thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return emptyList; } }); scheduler.onContainersAllocated(containers); drainableAppCallback.drain(); // no new allocation verify(mockApp, times(3)).taskAllocated(any(), any(), (Container) any()); // verify blacklisted container released verify(mockRMClient).releaseAssignedContainer(mockCId5); verify(mockRMClient, times(4)).releaseAssignedContainer((ContainerId) any()); // verify request added back verify(mockRMClient, times(6)).addContainerRequest(requestCaptor.capture()); CookieContainerRequest request5 = requestCaptor.getValue(); anyContainers.clear(); anyContainers.add(request5); Container mockContainer6 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer6.getNodeId().getHost()).thenReturn("host7"); ContainerId mockCId6 = mock(ContainerId.class); when(mockContainer6.getId()).thenReturn(mockCId6); containers.clear(); containers.add(mockContainer6); when(mockRMClient.getMatchingRequests((Priority) any(), eq(ResourceRequest.ANY), (Resource) any())) .thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return anyList; } }).thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return emptyList; } }); scheduler.onContainersAllocated(containers); drainableAppCallback.drain(); // new allocation verify(mockApp, times(4)).taskAllocated(any(), any(), (Container) any()); verify(mockApp).taskAllocated(mockTask4, mockCookie4, mockContainer6); // deallocate allocated task assertTrue(scheduler.deallocateTask(mockTask4, true)); drainableAppCallback.drain(); verify(mockApp).containerBeingReleased(mockCId6); verify(mockRMClient).releaseAssignedContainer(mockCId6); verify(mockRMClient, times(5)).releaseAssignedContainer((ContainerId) any()); // test unblacklist scheduler.unblacklistNode(badNodeId); verify(mockRMClient, times(1)).removeNodeFromBlacklist(badNodeId); assertEquals(0, scheduler.blacklistedNodes.size()); float progress = 0.5f; when(mockApp.getProgress()).thenReturn(progress); Assert.assertEquals(progress, scheduler.getProgress(), 0); // check duplicate allocation request scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1); drainableAppCallback.drain(); verify(mockRMClient, times(7)).addContainerRequest((CookieContainerRequest) any()); verify(mockRMClient, times(6)).removeContainerRequest((CookieContainerRequest) any()); scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1); drainableAppCallback.drain(); // old request removed and new one added verify(mockRMClient, times(7)).removeContainerRequest((CookieContainerRequest) any()); verify(mockRMClient, times(8)).addContainerRequest((CookieContainerRequest) any()); assertFalse(scheduler.deallocateTask(mockTask1, true)); List<NodeReport> mockUpdatedNodes = mock(List.class); scheduler.onNodesUpdated(mockUpdatedNodes); drainableAppCallback.drain(); verify(mockApp).nodesUpdated(mockUpdatedNodes); Exception mockException = mock(Exception.class); scheduler.onError(mockException); drainableAppCallback.drain(); verify(mockApp).onError(mockException); scheduler.onShutdownRequest(); drainableAppCallback.drain(); verify(mockApp).appShutdownRequested(); String appMsg = "success"; AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl); when(mockApp.getFinalAppStatus()).thenReturn(finalStatus); scheduler.stop(); drainableAppCallback.drain(); verify(mockRMClient).unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl); verify(mockRMClient).stop(); scheduler.close(); }
From source file:org.apache.tez.dag.app.rm.TestTaskScheduler.java
License:Apache License
@SuppressWarnings({ "unchecked" }) @Test(timeout = 10000)//from w w w. ja v a 2 s . com public void testTaskSchedulerWithReuse() throws Exception { RackResolver.init(new YarnConfiguration()); TaskSchedulerAppCallback mockApp = mock(TaskSchedulerAppCallback.class); AppContext mockAppContext = mock(AppContext.class); when(mockAppContext.getAMState()).thenReturn(DAGAppMasterState.RUNNING); TezAMRMClientAsync<CookieContainerRequest> mockRMClient = mock(TezAMRMClientAsync.class); String appHost = "host"; int appPort = 0; String appUrl = "url"; TaskSchedulerWithDrainableAppCallback scheduler = new TaskSchedulerWithDrainableAppCallback(mockApp, new AlwaysMatchesContainerMatcher(), appHost, appPort, appUrl, mockRMClient, mockAppContext); final TaskSchedulerAppCallbackDrainable drainableAppCallback = scheduler.getDrainableAppCallback(); Configuration conf = new Configuration(); // to match all in the same pass conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 0); // to release immediately after deallocate conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 0); scheduler.init(conf); drainableAppCallback.drain(); RegisterApplicationMasterResponse mockRegResponse = mock(RegisterApplicationMasterResponse.class); Resource mockMaxResource = mock(Resource.class); Map<ApplicationAccessType, String> mockAcls = mock(Map.class); when(mockRegResponse.getMaximumResourceCapability()).thenReturn(mockMaxResource); when(mockRegResponse.getApplicationACLs()).thenReturn(mockAcls); when(mockRMClient.registerApplicationMaster(anyString(), anyInt(), anyString())) .thenReturn(mockRegResponse); Resource mockClusterResource = mock(Resource.class); when(mockRMClient.getAvailableResources()).thenReturn(mockClusterResource); scheduler.start(); drainableAppCallback.drain(); Object mockTask1 = mock(Object.class); when(mockTask1.toString()).thenReturn("task1"); Object mockCookie1 = mock(Object.class); Resource mockCapability = mock(Resource.class); String[] hosts = { "host1", "host5" }; String[] racks = { "/default-rack", "/default-rack" }; final Priority mockPriority1 = Priority.newInstance(1); final Priority mockPriority2 = Priority.newInstance(2); final Priority mockPriority3 = Priority.newInstance(3); final Priority mockPriority4 = Priority.newInstance(4); final Priority mockPriority5 = Priority.newInstance(5); Object mockTask2 = mock(Object.class); when(mockTask2.toString()).thenReturn("task2"); Object mockCookie2 = mock(Object.class); Object mockTask3 = mock(Object.class); when(mockTask3.toString()).thenReturn("task3"); Object mockCookie3 = mock(Object.class); ArgumentCaptor<CookieContainerRequest> requestCaptor = ArgumentCaptor .forClass(CookieContainerRequest.class); scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority1, null, mockCookie1); drainableAppCallback.drain(); verify(mockRMClient, times(1)).addContainerRequest(requestCaptor.capture()); CookieContainerRequest request1 = requestCaptor.getValue(); scheduler.allocateTask(mockTask2, mockCapability, hosts, racks, mockPriority2, null, mockCookie2); drainableAppCallback.drain(); verify(mockRMClient, times(2)).addContainerRequest(requestCaptor.capture()); CookieContainerRequest request2 = requestCaptor.getValue(); scheduler.allocateTask(mockTask3, mockCapability, hosts, racks, mockPriority3, null, mockCookie3); drainableAppCallback.drain(); verify(mockRMClient, times(3)).addContainerRequest(requestCaptor.capture()); CookieContainerRequest request3 = requestCaptor.getValue(); List<Container> containers = new ArrayList<Container>(); // sending lower priority container first to make sure its not matched Container mockContainer4 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer4.getNodeId().getHost()).thenReturn("host4"); when(mockContainer4.toString()).thenReturn("container4"); when(mockContainer4.getPriority()).thenReturn(mockPriority4); ContainerId mockCId4 = mock(ContainerId.class); when(mockContainer4.getId()).thenReturn(mockCId4); when(mockCId4.toString()).thenReturn("container4"); containers.add(mockContainer4); Container mockContainer1 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer1.getNodeId().getHost()).thenReturn("host1"); when(mockContainer1.getPriority()).thenReturn(mockPriority1); when(mockContainer1.toString()).thenReturn("container1"); ContainerId mockCId1 = mock(ContainerId.class); when(mockContainer1.getId()).thenReturn(mockCId1); when(mockCId1.toString()).thenReturn("container1"); containers.add(mockContainer1); Container mockContainer2 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer2.getNodeId().getHost()).thenReturn("host2"); when(mockContainer2.getPriority()).thenReturn(mockPriority2); when(mockContainer2.toString()).thenReturn("container2"); ContainerId mockCId2 = mock(ContainerId.class); when(mockContainer2.getId()).thenReturn(mockCId2); when(mockCId2.toString()).thenReturn("container2"); containers.add(mockContainer2); Container mockContainer3 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer3.getNodeId().getHost()).thenReturn("host3"); when(mockContainer3.getPriority()).thenReturn(mockPriority3); when(mockContainer3.toString()).thenReturn("container3"); ContainerId mockCId3 = mock(ContainerId.class); when(mockContainer3.getId()).thenReturn(mockCId3); when(mockCId3.toString()).thenReturn("container3"); containers.add(mockContainer3); ArrayList<CookieContainerRequest> hostContainers = new ArrayList<CookieContainerRequest>(); hostContainers.add(request1); ArrayList<CookieContainerRequest> rackContainers = new ArrayList<CookieContainerRequest>(); rackContainers.add(request2); ArrayList<CookieContainerRequest> anyContainers = new ArrayList<CookieContainerRequest>(); anyContainers.add(request3); final List<ArrayList<CookieContainerRequest>> hostList = new LinkedList<ArrayList<CookieContainerRequest>>(); hostList.add(hostContainers); final List<ArrayList<CookieContainerRequest>> rackList = new LinkedList<ArrayList<CookieContainerRequest>>(); rackList.add(rackContainers); final List<ArrayList<CookieContainerRequest>> anyList = new LinkedList<ArrayList<CookieContainerRequest>>(); anyList.add(anyContainers); final List<ArrayList<CookieContainerRequest>> emptyList = new LinkedList<ArrayList<CookieContainerRequest>>(); // return pri1 requests for host1 when(mockRMClient.getMatchingRequestsForTopPriority(eq("host1"), (Resource) any())) .thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return hostList; } }); // second request matched to rack. RackResolver by default puts hosts in // /default-rack. We need to workaround by returning rack matches only once when(mockRMClient.getMatchingRequestsForTopPriority(eq("/default-rack"), (Resource) any())) .thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return rackList; } }).thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return emptyList; } }); // third request matched to ANY when(mockRMClient.getMatchingRequestsForTopPriority(eq(ResourceRequest.ANY), (Resource) any())) .thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return anyList; } }).thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return emptyList; } }); when(mockRMClient.getTopPriority()).then(new Answer<Priority>() { @Override public Priority answer(InvocationOnMock invocation) throws Throwable { int allocations = drainableAppCallback.count.get(); if (allocations == 0) { return mockPriority1; } if (allocations == 1) { return mockPriority2; } if (allocations == 2) { return mockPriority3; } if (allocations == 3) { return mockPriority4; } return null; } }); AtomicBoolean drainNotifier = new AtomicBoolean(false); scheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier; scheduler.onContainersAllocated(containers); TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier); drainableAppCallback.drain(); // exact number allocations returned verify(mockApp, times(3)).taskAllocated(any(), any(), (Container) any()); // first container allocated verify(mockApp).taskAllocated(mockTask1, mockCookie1, mockContainer1); verify(mockApp).taskAllocated(mockTask2, mockCookie2, mockContainer2); verify(mockApp).taskAllocated(mockTask3, mockCookie3, mockContainer3); verify(mockRMClient).removeContainerRequest(request1); verify(mockRMClient).removeContainerRequest(request2); verify(mockRMClient).removeContainerRequest(request3); // verify unwanted container released verify(mockRMClient).releaseAssignedContainer(mockCId4); // deallocate allocated task assertTrue(scheduler.deallocateTask(mockTask1, true)); drainableAppCallback.drain(); verify(mockApp).containerBeingReleased(mockCId1); verify(mockRMClient).releaseAssignedContainer(mockCId1); // deallocate allocated container Assert.assertEquals(mockTask2, scheduler.deallocateContainer(mockCId2)); drainableAppCallback.drain(); verify(mockRMClient).releaseAssignedContainer(mockCId2); verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any()); List<ContainerStatus> statuses = new ArrayList<ContainerStatus>(); ContainerStatus mockStatus1 = mock(ContainerStatus.class); when(mockStatus1.getContainerId()).thenReturn(mockCId1); statuses.add(mockStatus1); ContainerStatus mockStatus2 = mock(ContainerStatus.class); when(mockStatus2.getContainerId()).thenReturn(mockCId2); statuses.add(mockStatus2); ContainerStatus mockStatus3 = mock(ContainerStatus.class); when(mockStatus3.getContainerId()).thenReturn(mockCId3); statuses.add(mockStatus3); ContainerStatus mockStatus4 = mock(ContainerStatus.class); when(mockStatus4.getContainerId()).thenReturn(mockCId4); statuses.add(mockStatus4); scheduler.onContainersCompleted(statuses); drainableAppCallback.drain(); // released container status returned verify(mockApp).containerCompleted(mockTask1, mockStatus1); verify(mockApp).containerCompleted(mockTask2, mockStatus2); // currently allocated container status returned and not released verify(mockApp).containerCompleted(mockTask3, mockStatus3); // no other statuses returned verify(mockApp, times(3)).containerCompleted(any(), (ContainerStatus) any()); verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any()); // verify blacklisting verify(mockRMClient, times(0)).addNodeToBlacklist((NodeId) any()); String badHost = "host6"; NodeId badNodeId = mock(NodeId.class); when(badNodeId.getHost()).thenReturn(badHost); scheduler.blacklistNode(badNodeId); verify(mockRMClient, times(1)).addNodeToBlacklist(badNodeId); Object mockTask4 = mock(Object.class); when(mockTask4.toString()).thenReturn("task4"); Object mockCookie4 = mock(Object.class); scheduler.allocateTask(mockTask4, mockCapability, null, null, mockPriority4, null, mockCookie4); drainableAppCallback.drain(); verify(mockRMClient, times(4)).addContainerRequest(requestCaptor.capture()); CookieContainerRequest request4 = requestCaptor.getValue(); anyContainers.clear(); anyContainers.add(request4); Container mockContainer5 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer5.getNodeId().getHost()).thenReturn(badHost); when(mockContainer5.getNodeId()).thenReturn(badNodeId); ContainerId mockCId5 = mock(ContainerId.class); when(mockContainer5.toString()).thenReturn("container5"); when(mockCId5.toString()).thenReturn("container5"); when(mockContainer5.getId()).thenReturn(mockCId5); when(mockContainer5.getPriority()).thenReturn(mockPriority4); containers.clear(); containers.add(mockContainer5); when(mockRMClient.getMatchingRequestsForTopPriority(eq(ResourceRequest.ANY), (Resource) any())) .thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return anyList; } }).thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return emptyList; } }); drainNotifier.set(false); scheduler.onContainersAllocated(containers); TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier); drainableAppCallback.drain(); // no new allocation verify(mockApp, times(3)).taskAllocated(any(), any(), (Container) any()); // verify blacklisted container released verify(mockRMClient).releaseAssignedContainer(mockCId5); verify(mockRMClient, times(4)).releaseAssignedContainer((ContainerId) any()); // verify request added back verify(mockRMClient, times(5)).addContainerRequest(requestCaptor.capture()); CookieContainerRequest request5 = requestCaptor.getValue(); anyContainers.clear(); anyContainers.add(request5); Container mockContainer6 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer6.getNodeId().getHost()).thenReturn("host7"); ContainerId mockCId6 = mock(ContainerId.class); when(mockContainer6.getId()).thenReturn(mockCId6); when(mockContainer6.toString()).thenReturn("container6"); when(mockCId6.toString()).thenReturn("container6"); containers.clear(); containers.add(mockContainer6); when(mockRMClient.getMatchingRequestsForTopPriority(eq(ResourceRequest.ANY), (Resource) any())) .thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return anyList; } }).thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return emptyList; } }); drainNotifier.set(false); scheduler.onContainersAllocated(containers); TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier); drainableAppCallback.drain(); // new allocation verify(mockApp, times(4)).taskAllocated(any(), any(), (Container) any()); verify(mockApp).taskAllocated(mockTask4, mockCookie4, mockContainer6); // deallocate allocated task assertTrue(scheduler.deallocateTask(mockTask4, true)); drainableAppCallback.drain(); verify(mockApp).containerBeingReleased(mockCId6); verify(mockRMClient).releaseAssignedContainer(mockCId6); verify(mockRMClient, times(5)).releaseAssignedContainer((ContainerId) any()); // test unblacklist scheduler.unblacklistNode(badNodeId); verify(mockRMClient, times(1)).removeNodeFromBlacklist(badNodeId); assertEquals(0, scheduler.blacklistedNodes.size()); // verify container level matching // fudge the top level priority to prevent containers from being released // if top level priority is higher than newly allocated containers then // they will not be released final AtomicBoolean fudgePriority = new AtomicBoolean(true); when(mockRMClient.getTopPriority()).then(new Answer<Priority>() { @Override public Priority answer(InvocationOnMock invocation) throws Throwable { if (fudgePriority.get()) { return mockPriority4; } return mockPriority5; } }); // add a dummy task to prevent release of allocated containers Object mockTask5 = mock(Object.class); when(mockTask5.toString()).thenReturn("task5"); Object mockCookie5 = mock(Object.class); scheduler.allocateTask(mockTask5, mockCapability, hosts, racks, mockPriority5, null, mockCookie5); verify(mockRMClient, times(6)).addContainerRequest(requestCaptor.capture()); CookieContainerRequest request6 = requestCaptor.getValue(); drainableAppCallback.drain(); // add containers so that we can reference one of them for container specific // allocation containers.clear(); Container mockContainer7 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer7.getNodeId().getHost()).thenReturn("host5"); ContainerId mockCId7 = mock(ContainerId.class); when(mockContainer7.toString()).thenReturn("container7"); when(mockCId7.toString()).thenReturn("container7"); when(mockContainer7.getId()).thenReturn(mockCId7); when(mockContainer7.getPriority()).thenReturn(mockPriority5); containers.add(mockContainer7); Container mockContainer8 = mock(Container.class, RETURNS_DEEP_STUBS); when(mockContainer8.getNodeId().getHost()).thenReturn("host5"); ContainerId mockCId8 = mock(ContainerId.class); when(mockContainer8.toString()).thenReturn("container8"); when(mockCId8.toString()).thenReturn("container8"); when(mockContainer8.getId()).thenReturn(mockCId8); when(mockContainer8.getPriority()).thenReturn(mockPriority5); containers.add(mockContainer8); drainNotifier.set(false); scheduler.onContainersAllocated(containers); drainableAppCallback.drain(); verify(mockRMClient, times(5)).releaseAssignedContainer((ContainerId) any()); Object mockTask6 = mock(Object.class); when(mockTask6.toString()).thenReturn("task6"); Object mockCookie6 = mock(Object.class); // allocate request with container affinity scheduler.allocateTask(mockTask6, mockCapability, mockCId7, mockPriority5, null, mockCookie6); drainableAppCallback.drain(); verify(mockRMClient, times(7)).addContainerRequest(requestCaptor.capture()); CookieContainerRequest request7 = requestCaptor.getValue(); hostContainers.clear(); hostContainers.add(request6); hostContainers.add(request7); when(mockRMClient.getMatchingRequestsForTopPriority(eq("host5"), (Resource) any())) .thenAnswer(new Answer<List<? extends Collection<CookieContainerRequest>>>() { @Override public List<? extends Collection<CookieContainerRequest>> answer(InvocationOnMock invocation) throws Throwable { return hostList; } }); // stop fudging top priority fudgePriority.set(false); TestTaskSchedulerHelpers.waitForDelayedDrainNotify(drainNotifier); drainableAppCallback.drain(); verify(mockApp, times(6)).taskAllocated(any(), any(), (Container) any()); // container7 allocated to the task with affinity for it verify(mockApp).taskAllocated(mockTask6, mockCookie6, mockContainer7); // deallocate allocated task assertTrue(scheduler.deallocateTask(mockTask5, true)); assertTrue(scheduler.deallocateTask(mockTask6, true)); drainableAppCallback.drain(); verify(mockApp).containerBeingReleased(mockCId7); verify(mockApp).containerBeingReleased(mockCId8); verify(mockRMClient).releaseAssignedContainer(mockCId7); verify(mockRMClient).releaseAssignedContainer(mockCId8); verify(mockRMClient, times(7)).releaseAssignedContainer((ContainerId) any()); float progress = 0.5f; when(mockApp.getProgress()).thenReturn(progress); Assert.assertEquals(progress, scheduler.getProgress(), 0); List<NodeReport> mockUpdatedNodes = mock(List.class); scheduler.onNodesUpdated(mockUpdatedNodes); drainableAppCallback.drain(); verify(mockApp).nodesUpdated(mockUpdatedNodes); Exception mockException = mock(Exception.class); scheduler.onError(mockException); drainableAppCallback.drain(); verify(mockApp).onError(mockException); scheduler.onShutdownRequest(); drainableAppCallback.drain(); verify(mockApp).appShutdownRequested(); String appMsg = "success"; AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl); when(mockApp.getFinalAppStatus()).thenReturn(finalStatus); scheduler.stop(); drainableAppCallback.drain(); verify(mockRMClient).unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl); verify(mockRMClient).stop(); scheduler.close(); }
From source file:org.apache.tez.dag.app.rm.TezAMRMClientAsync.java
License:Apache License
public synchronized void addNodeToBlacklist(NodeId nodeId) { client.updateBlacklist(Collections.singletonList(nodeId.getHost()), null); }
From source file:org.apache.tez.dag.app.rm.TezAMRMClientAsync.java
License:Apache License
public synchronized void removeNodeFromBlacklist(NodeId nodeId) { client.updateBlacklist(null, Collections.singletonList(nodeId.getHost())); }
From source file:org.apache.tez.dag.app.TaskCommunicatorManager.java
License:Apache License
@Override public void registerRunningContainer(ContainerId containerId, int taskCommId) { if (LOG.isDebugEnabled()) { LOG.debug("ContainerId: " + containerId + " registered with TaskAttemptListener"); }// w w w .j ava2 s . c o m ContainerInfo oldInfo = registeredContainers.put(containerId, NULL_CONTAINER_INFO); if (oldInfo != null) { throw new TezUncheckedException("Multiple registrations for containerId: " + containerId); } NodeId nodeId = context.getAllContainers().get(containerId).getContainer().getNodeId(); try { taskCommunicators[taskCommId].registerRunningContainer(containerId, nodeId.getHost(), nodeId.getPort()); } catch (Exception e) { String msg = "Error in TaskCommunicator when registering running Container" + ", communicator=" + Utils.getTaskCommIdentifierString(taskCommId, context) + ", containerId=" + containerId + ", nodeId=" + nodeId; LOG.error(msg, e); sendEvent(new DAGAppMasterEventUserServiceFatalError( DAGAppMasterEventType.TASK_COMMUNICATOR_SERVICE_FATAL_ERROR, msg, e)); } }