/** * 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; import java.util.ArrayList; import java.util.List; import java.util.Set; import io.hops.util.DBUtility; import io.hops.util.RMStorageFactory; import io.hops.util.YarnAPIStorageFactory; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.security.SecurityUtilTestHelper; import org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse; import org.apache.hadoop.yarn.api.records.ApplicationAttemptId; import org.apache.hadoop.yarn.api.records.Container; import org.apache.hadoop.yarn.api.records.ContainerId; import org.apache.hadoop.yarn.api.records.LogAggregationContext; import org.apache.hadoop.yarn.api.records.NodeId; import org.apache.hadoop.yarn.api.records.Priority; import org.apache.hadoop.yarn.api.records.Resource; import org.apache.hadoop.yarn.api.records.ResourceRequest; import org.apache.hadoop.yarn.api.records.Token; import org.apache.hadoop.yarn.conf.YarnConfiguration; import org.apache.hadoop.yarn.security.ContainerTokenIdentifier; import org.apache.hadoop.yarn.server.resourcemanager.MockAM; import org.apache.hadoop.yarn.server.resourcemanager.MockNM; import org.apache.hadoop.yarn.server.resourcemanager.MockRM; import org.apache.hadoop.yarn.server.resourcemanager.RMContextImpl; import org.apache.hadoop.yarn.server.resourcemanager.RMSecretManagerService; import org.apache.hadoop.yarn.server.resourcemanager.ResourceManager; import org.apache.hadoop.yarn.server.resourcemanager.TestFifoScheduler; import org.apache.hadoop.yarn.server.resourcemanager.nodelabels.NullRMNodeLabelsManager; import org.apache.hadoop.yarn.server.resourcemanager.nodelabels.RMNodeLabelsManager; import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMApp; import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttempt; import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttemptState; import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainer; import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainerState; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.ResourceScheduler; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.SchedulerAppReport; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.YarnScheduler; import org.apache.hadoop.yarn.server.resourcemanager.security.RMContainerTokenSecretManager; import org.apache.hadoop.yarn.server.utils.BuilderUtils; import org.junit.Assert; import org.junit.Before; import org.junit.Test; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Sets; public class TestContainerAllocation { private static final Log LOG = LogFactory.getLog(TestFifoScheduler.class); private final int GB = 1024; private YarnConfiguration conf; RMNodeLabelsManager mgr; @Before public void setUp() throws Exception { conf = new YarnConfiguration(); conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class); RMStorageFactory.setConfiguration(conf); YarnAPIStorageFactory.setConfiguration(conf); DBUtility.InitializeDB(); mgr = new NullRMNodeLabelsManager(); mgr.init(conf); } @Test(timeout = 3000000) public void testExcessReservationThanNodeManagerCapacity() throws Exception { @SuppressWarnings("resource") MockRM rm = new MockRM(conf); rm.start(); // Register node1 MockNM nm1 = rm.registerNode("127.0.0.1:1234", 2 * GB, 4); MockNM nm2 = rm.registerNode("127.0.0.1:2234", 3 * GB, 4); nm1.nodeHeartbeat(true); nm2.nodeHeartbeat(true); // wait.. int waitCount = 20; int size = rm.getRMContext().getRMNodes().size(); while ((size = rm.getRMContext().getRMNodes().size()) != 2 && waitCount-- > 0) { LOG.info("Waiting for node managers to register : " + size); Thread.sleep(100); } Assert.assertEquals(2, rm.getRMContext().getRMNodes().size()); // Submit an application RMApp app1 = rm.submitApp(128); // kick the scheduling nm1.nodeHeartbeat(true); RMAppAttempt attempt1 = app1.getCurrentAppAttempt(); MockAM am1 = rm.sendAMLaunched(attempt1.getAppAttemptId()); am1.registerAppAttempt(); LOG.info("sending container requests "); am1.addRequests(new String[] {"*"}, 2 * GB, 1, 1); AllocateResponse alloc1Response = am1.schedule(); // send the request // kick the scheduler nm1.nodeHeartbeat(true); int waitCounter = 20; LOG.info("heartbeating nm1"); while (alloc1Response.getAllocatedContainers().size() < 1 && waitCounter-- > 0) { LOG.info("Waiting for containers to be created for app 1..."); Thread.sleep(500); alloc1Response = am1.schedule(); } LOG.info("received container : " + alloc1Response.getAllocatedContainers().size()); // No container should be allocated. // Internally it should not been reserved. Assert.assertTrue(alloc1Response.getAllocatedContainers().size() == 0); LOG.info("heartbeating nm2"); waitCounter = 20; nm2.nodeHeartbeat(true); while (alloc1Response.getAllocatedContainers().size() < 1 && waitCounter-- > 0) { LOG.info("Waiting for containers to be created for app 1..."); Thread.sleep(500); alloc1Response = am1.schedule(); } LOG.info("received container : " + alloc1Response.getAllocatedContainers().size()); Assert.assertTrue(alloc1Response.getAllocatedContainers().size() == 1); rm.stop(); } // This is to test container tokens are generated when the containers are // acquired by the AM, not when the containers are allocated @Test public void testContainerTokenGeneratedOnPullRequest() throws Exception { MockRM rm1 = new MockRM(conf); rm1.start(); MockNM nm1 = rm1.registerNode("127.0.0.1:1234", 8000); RMApp app1 = rm1.submitApp(200); MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1); // request a container. am1.allocate("127.0.0.1", 1024, 1, new ArrayList<ContainerId>()); ContainerId containerId2 = ContainerId.newContainerId(am1.getApplicationAttemptId(), 2); rm1.waitForState(nm1, containerId2, RMContainerState.ALLOCATED); RMContainer container = rm1.getResourceScheduler().getRMContainer(containerId2); // no container token is generated. Assert.assertEquals(containerId2, container.getContainerId()); Assert.assertNull(container.getContainer().getContainerToken()); // acquire the container. List<Container> containers = am1.allocate(new ArrayList<ResourceRequest>(), new ArrayList<ContainerId>()).getAllocatedContainers(); Assert.assertEquals(containerId2, containers.get(0).getId()); // container token is generated. Assert.assertNotNull(containers.get(0).getContainerToken()); rm1.stop(); } @Test public void testNormalContainerAllocationWhenDNSUnavailable() throws Exception{ MockRM rm1 = new MockRM(conf); rm1.start(); MockNM nm1 = rm1.registerNode("unknownhost:1234", 8000); RMApp app1 = rm1.submitApp(200); MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1); // request a container. am1.allocate("127.0.0.1", 1024, 1, new ArrayList<ContainerId>()); ContainerId containerId2 = ContainerId.newContainerId(am1.getApplicationAttemptId(), 2); rm1.waitForState(nm1, containerId2, RMContainerState.ALLOCATED); // acquire the container. SecurityUtilTestHelper.setTokenServiceUseIp(true); List<Container> containers = am1.allocate(new ArrayList<ResourceRequest>(), new ArrayList<ContainerId>()).getAllocatedContainers(); // not able to fetch the container; Assert.assertEquals(0, containers.size()); SecurityUtilTestHelper.setTokenServiceUseIp(false); containers = am1.allocate(new ArrayList<ResourceRequest>(), new ArrayList<ContainerId>()).getAllocatedContainers(); // should be able to fetch the container; Assert.assertEquals(1, containers.size()); } // This is to test whether LogAggregationContext is passed into // container tokens correctly @Test public void testLogAggregationContextPassedIntoContainerToken() throws Exception { MockRM rm1 = new MockRM(conf); rm1.start(); MockNM nm1 = rm1.registerNode("127.0.0.1:1234", 8000); MockNM nm2 = rm1.registerNode("127.0.0.1:2345", 8000); // LogAggregationContext is set as null Assert .assertNull(getLogAggregationContextFromContainerToken(rm1, nm1, null)); // create a not-null LogAggregationContext LogAggregationContext logAggregationContext = LogAggregationContext.newInstance( "includePattern", "excludePattern", "rolledLogsIncludePattern", "rolledLogsExcludePattern"); LogAggregationContext returned = getLogAggregationContextFromContainerToken(rm1, nm2, logAggregationContext); Assert.assertEquals("includePattern", returned.getIncludePattern()); Assert.assertEquals("excludePattern", returned.getExcludePattern()); Assert.assertEquals("rolledLogsIncludePattern", returned.getRolledLogsIncludePattern()); Assert.assertEquals("rolledLogsExcludePattern", returned.getRolledLogsExcludePattern()); rm1.stop(); } private LogAggregationContext getLogAggregationContextFromContainerToken( MockRM rm1, MockNM nm1, LogAggregationContext logAggregationContext) throws Exception { RMApp app2 = rm1.submitApp(200, logAggregationContext); MockAM am2 = MockRM.launchAndRegisterAM(app2, rm1, nm1); nm1.nodeHeartbeat(true); // request a container. am2.allocate("127.0.0.1", 512, 1, new ArrayList<ContainerId>()); ContainerId containerId = ContainerId.newContainerId(am2.getApplicationAttemptId(), 2); rm1.waitForState(nm1, containerId, RMContainerState.ALLOCATED); // acquire the container. List<Container> containers = am2.allocate(new ArrayList<ResourceRequest>(), new ArrayList<ContainerId>()).getAllocatedContainers(); Assert.assertEquals(containerId, containers.get(0).getId()); // container token is generated. Assert.assertNotNull(containers.get(0).getContainerToken()); ContainerTokenIdentifier token = BuilderUtils.newContainerTokenIdentifier(containers.get(0) .getContainerToken()); return token.getLogAggregationContext(); } private volatile int numRetries = 0; private class TestRMSecretManagerService extends RMSecretManagerService { public TestRMSecretManagerService(Configuration conf, RMContextImpl rmContext) { super(conf, rmContext); } @Override protected RMContainerTokenSecretManager createContainerTokenSecretManager( Configuration conf) { return new RMContainerTokenSecretManager(conf) { @Override public Token createContainerToken(ContainerId containerId, NodeId nodeId, String appSubmitter, Resource capability, Priority priority, long createTime, LogAggregationContext logAggregationContext) { numRetries++; return super.createContainerToken(containerId, nodeId, appSubmitter, capability, priority, createTime, logAggregationContext); } }; } } // This is to test fetching AM container will be retried, if AM container is // not fetchable since DNS is unavailable causing container token/NMtoken // creation failure. @Test(timeout = 30000) public void testAMContainerAllocationWhenDNSUnavailable() throws Exception { MockRM rm1 = new MockRM(conf) { @Override protected RMSecretManagerService createRMSecretManagerService() { return new TestRMSecretManagerService(conf, rmContext); } }; rm1.start(); MockNM nm1 = rm1.registerNode("unknownhost:1234", 8000); SecurityUtilTestHelper.setTokenServiceUseIp(true); RMApp app1 = rm1.submitApp(200); RMAppAttempt attempt = app1.getCurrentAppAttempt(); nm1.nodeHeartbeat(true); // fetching am container will fail, keep retrying 5 times. while (numRetries <= 5) { nm1.nodeHeartbeat(true); Thread.sleep(1000); Assert.assertEquals(RMAppAttemptState.SCHEDULED, attempt.getAppAttemptState()); System.out.println("Waiting for am container to be allocated."); } SecurityUtilTestHelper.setTokenServiceUseIp(false); rm1.waitForState(attempt.getAppAttemptId(), RMAppAttemptState.ALLOCATED); MockRM.launchAndRegisterAM(app1, rm1, nm1); } private Configuration getConfigurationWithDefaultQueueLabels( Configuration config) { final String A = CapacitySchedulerConfiguration.ROOT + ".a"; final String B = CapacitySchedulerConfiguration.ROOT + ".b"; CapacitySchedulerConfiguration conf = (CapacitySchedulerConfiguration) getConfigurationWithQueueLabels(config); new CapacitySchedulerConfiguration(config); conf.setDefaultNodeLabelExpression(A, "x"); conf.setDefaultNodeLabelExpression(B, "y"); return conf; } private Configuration getConfigurationWithQueueLabels(Configuration config) { CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration(config); // Define top-level queues conf.setQueues(CapacitySchedulerConfiguration.ROOT, new String[] {"a", "b", "c"}); conf.setCapacityByLabel(CapacitySchedulerConfiguration.ROOT, "x", 100); conf.setCapacityByLabel(CapacitySchedulerConfiguration.ROOT, "y", 100); final String A = CapacitySchedulerConfiguration.ROOT + ".a"; conf.setCapacity(A, 10); conf.setMaximumCapacity(A, 15); conf.setAccessibleNodeLabels(A, toSet("x")); conf.setCapacityByLabel(A, "x", 100); final String B = CapacitySchedulerConfiguration.ROOT + ".b"; conf.setCapacity(B, 20); conf.setAccessibleNodeLabels(B, toSet("y")); conf.setCapacityByLabel(B, "y", 100); final String C = CapacitySchedulerConfiguration.ROOT + ".c"; conf.setCapacity(C, 70); conf.setMaximumCapacity(C, 70); conf.setAccessibleNodeLabels(C, RMNodeLabelsManager.EMPTY_STRING_SET); // Define 2nd-level queues final String A1 = A + ".a1"; conf.setQueues(A, new String[] {"a1"}); conf.setCapacity(A1, 100); conf.setMaximumCapacity(A1, 100); conf.setCapacityByLabel(A1, "x", 100); final String B1 = B + ".b1"; conf.setQueues(B, new String[] {"b1"}); conf.setCapacity(B1, 100); conf.setMaximumCapacity(B1, 100); conf.setCapacityByLabel(B1, "y", 100); final String C1 = C + ".c1"; conf.setQueues(C, new String[] {"c1"}); conf.setCapacity(C1, 100); conf.setMaximumCapacity(C1, 100); return conf; } private void checkTaskContainersHost(ApplicationAttemptId attemptId, ContainerId containerId, ResourceManager rm, String host) { YarnScheduler scheduler = rm.getRMContext().getScheduler(); SchedulerAppReport appReport = scheduler.getSchedulerAppInfo(attemptId); Assert.assertTrue(appReport.getLiveContainers().size() > 0); for (RMContainer c : appReport.getLiveContainers()) { if (c.getContainerId().equals(containerId)) { Assert.assertEquals(host, c.getAllocatedNode().getHost()); } } } @SuppressWarnings("unchecked") private <E> Set<E> toSet(E... elements) { Set<E> set = Sets.newHashSet(elements); return set; } private Configuration getComplexConfigurationWithQueueLabels( Configuration config) { CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration(config); // Define top-level queues conf.setQueues(CapacitySchedulerConfiguration.ROOT, new String[] {"a", "b"}); conf.setCapacityByLabel(CapacitySchedulerConfiguration.ROOT, "x", 100); conf.setCapacityByLabel(CapacitySchedulerConfiguration.ROOT, "y", 100); conf.setCapacityByLabel(CapacitySchedulerConfiguration.ROOT, "z", 100); final String A = CapacitySchedulerConfiguration.ROOT + ".a"; conf.setCapacity(A, 10); conf.setMaximumCapacity(A, 10); conf.setAccessibleNodeLabels(A, toSet("x", "y")); conf.setCapacityByLabel(A, "x", 100); conf.setCapacityByLabel(A, "y", 50); final String B = CapacitySchedulerConfiguration.ROOT + ".b"; conf.setCapacity(B, 90); conf.setMaximumCapacity(B, 100); conf.setAccessibleNodeLabels(B, toSet("y", "z")); conf.setCapacityByLabel(B, "y", 50); conf.setCapacityByLabel(B, "z", 100); // Define 2nd-level queues final String A1 = A + ".a1"; conf.setQueues(A, new String[] {"a1"}); conf.setCapacity(A1, 100); conf.setMaximumCapacity(A1, 100); conf.setAccessibleNodeLabels(A1, toSet("x", "y")); conf.setDefaultNodeLabelExpression(A1, "x"); conf.setCapacityByLabel(A1, "x", 100); conf.setCapacityByLabel(A1, "y", 100); conf.setQueues(B, new String[] {"b1", "b2"}); final String B1 = B + ".b1"; conf.setCapacity(B1, 50); conf.setMaximumCapacity(B1, 50); conf.setAccessibleNodeLabels(B1, RMNodeLabelsManager.EMPTY_STRING_SET); final String B2 = B + ".b2"; conf.setCapacity(B2, 50); conf.setMaximumCapacity(B2, 50); conf.setAccessibleNodeLabels(B2, toSet("y", "z")); conf.setCapacityByLabel(B2, "y", 100); conf.setCapacityByLabel(B2, "z", 100); return conf; } @Test (timeout = 300000) public void testContainerAllocationWithSingleUserLimits() throws Exception { final RMNodeLabelsManager mgr = new NullRMNodeLabelsManager(); mgr.init(conf); // set node -> label mgr.addToCluserNodeLabels(ImmutableSet.of("x", "y")); mgr.addLabelsToNode(ImmutableMap.of(NodeId.newInstance("h1", 0), toSet("x"), NodeId.newInstance("h2", 0), toSet("y"))); // inject node label manager MockRM rm1 = new MockRM(getConfigurationWithDefaultQueueLabels(conf)) { @Override public RMNodeLabelsManager createNodeLabelManager() { return mgr; } }; rm1.getRMContext().setNodeLabelManager(mgr); rm1.start(); MockNM nm1 = rm1.registerNode("h1:1234", 8000); // label = x rm1.registerNode("h2:1234", 8000); // label = y MockNM nm3 = rm1.registerNode("h3:1234", 8000); // label = <empty> // launch an app to queue a1 (label = x), and check all container will // be allocated in h1 RMApp app1 = rm1.submitApp(200, "app", "user", null, "a1"); MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1); // A has only 10% of x, so it can only allocate one container in label=empty ContainerId containerId = ContainerId.newContainerId(am1.getApplicationAttemptId(), 2); am1.allocate("*", 1024, 1, new ArrayList<ContainerId>(), ""); Assert.assertTrue(rm1.waitForState(nm3, containerId, RMContainerState.ALLOCATED, 10 * 1000)); // Cannot allocate 2nd label=empty container containerId = ContainerId.newContainerId(am1.getApplicationAttemptId(), 3); am1.allocate("*", 1024, 1, new ArrayList<ContainerId>(), ""); Assert.assertFalse(rm1.waitForState(nm3, containerId, RMContainerState.ALLOCATED, 10 * 1000)); // A has default user limit = 100, so it can use all resource in label = x // We can allocate floor(8000 / 1024) = 7 containers for (int id = 3; id <= 8; id++) { containerId = ContainerId.newContainerId(am1.getApplicationAttemptId(), id); am1.allocate("*", 1024, 1, new ArrayList<ContainerId>(), "x"); Assert.assertTrue(rm1.waitForState(nm1, containerId, RMContainerState.ALLOCATED, 10 * 1000)); } rm1.close(); } @Test(timeout = 300000) public void testContainerAllocateWithComplexLabels() throws Exception { /* * Queue structure: * root (*) * ________________ * / \ * a x(100%), y(50%) b y(50%), z(100%) * ________________ ______________ * / / \ * a1 (x,y) b1(no) b2(y,z) * 100% y = 100%, z = 100% * * Node structure: * h1 : x * h2 : y * h3 : y * h4 : z * h5 : NO * * Total resource: * x: 4G * y: 6G * z: 2G * *: 2G * * Resource of * a1: x=4G, y=3G, NO=0.2G * b1: NO=0.9G (max=1G) * b2: y=3, z=2G, NO=0.9G (max=1G) * * Each node can only allocate two containers */ // set node -> label mgr.addToCluserNodeLabels(ImmutableSet.of("x", "y", "z")); mgr.addLabelsToNode(ImmutableMap.of(NodeId.newInstance("h1", 0), toSet("x"), NodeId.newInstance("h2", 0), toSet("y"), NodeId.newInstance("h3", 0), toSet("y"), NodeId.newInstance("h4", 0), toSet("z"), NodeId.newInstance("h5", 0), RMNodeLabelsManager.EMPTY_STRING_SET)); // inject node label manager MockRM rm1 = new MockRM(getComplexConfigurationWithQueueLabels(conf)) { @Override public RMNodeLabelsManager createNodeLabelManager() { return mgr; } }; rm1.getRMContext().setNodeLabelManager(mgr); rm1.start(); MockNM nm1 = rm1.registerNode("h1:1234", 2048); MockNM nm2 = rm1.registerNode("h2:1234", 2048); MockNM nm3 = rm1.registerNode("h3:1234", 2048); MockNM nm4 = rm1.registerNode("h4:1234", 2048); MockNM nm5 = rm1.registerNode("h5:1234", 2048); ContainerId containerId; // launch an app to queue a1 (label = x), and check all container will // be allocated in h1 RMApp app1 = rm1.submitApp(1024, "app", "user", null, "a1"); MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1); // request a container (label = y). can be allocated on nm2 am1.allocate("*", 1024, 1, new ArrayList<ContainerId>(), "y"); containerId = ContainerId.newContainerId(am1.getApplicationAttemptId(), 2L); Assert.assertTrue(rm1.waitForState(nm2, containerId, RMContainerState.ALLOCATED, 10 * 1000)); checkTaskContainersHost(am1.getApplicationAttemptId(), containerId, rm1, "h2"); // launch an app to queue b1 (label = y), and check all container will // be allocated in h5 RMApp app2 = rm1.submitApp(1024, "app", "user", null, "b1"); MockAM am2 = MockRM.launchAndRegisterAM(app2, rm1, nm5); // request a container for AM, will succeed // and now b1's queue capacity will be used, cannot allocate more containers // (Maximum capacity reached) am2.allocate("*", 1024, 1, new ArrayList<ContainerId>()); containerId = ContainerId.newContainerId(am2.getApplicationAttemptId(), 2); Assert.assertFalse(rm1.waitForState(nm4, containerId, RMContainerState.ALLOCATED, 10 * 1000)); Assert.assertFalse(rm1.waitForState(nm5, containerId, RMContainerState.ALLOCATED, 10 * 1000)); // launch an app to queue b2 RMApp app3 = rm1.submitApp(1024, "app", "user", null, "b2"); MockAM am3 = MockRM.launchAndRegisterAM(app3, rm1, nm5); // request a container. try to allocate on nm1 (label = x) and nm3 (label = // y,z). Will successfully allocate on nm3 am3.allocate("*", 1024, 1, new ArrayList<ContainerId>(), "y"); containerId = ContainerId.newContainerId(am3.getApplicationAttemptId(), 2); Assert.assertFalse(rm1.waitForState(nm1, containerId, RMContainerState.ALLOCATED, 10 * 1000)); Assert.assertTrue(rm1.waitForState(nm3, containerId, RMContainerState.ALLOCATED, 10 * 1000)); checkTaskContainersHost(am3.getApplicationAttemptId(), containerId, rm1, "h3"); // try to allocate container (request label = z) on nm4 (label = y,z). // Will successfully allocate on nm4 only. am3.allocate("*", 1024, 1, new ArrayList<ContainerId>(), "z"); containerId = ContainerId.newContainerId(am3.getApplicationAttemptId(), 3L); Assert.assertTrue(rm1.waitForState(nm4, containerId, RMContainerState.ALLOCATED, 10 * 1000)); checkTaskContainersHost(am3.getApplicationAttemptId(), containerId, rm1, "h4"); rm1.close(); } @Test (timeout = 120000) public void testContainerAllocateWithLabels() throws Exception { // set node -> label mgr.addToCluserNodeLabels(ImmutableSet.of("x", "y")); mgr.addLabelsToNode(ImmutableMap.of(NodeId.newInstance("h1", 0), toSet("x"), NodeId.newInstance("h2", 0), toSet("y"))); // inject node label manager MockRM rm1 = new MockRM(getConfigurationWithQueueLabels(conf)) { @Override public RMNodeLabelsManager createNodeLabelManager() { return mgr; } }; rm1.getRMContext().setNodeLabelManager(mgr); rm1.start(); MockNM nm1 = rm1.registerNode("h1:1234", 8000); // label = x MockNM nm2 = rm1.registerNode("h2:1234", 8000); // label = y MockNM nm3 = rm1.registerNode("h3:1234", 8000); // label = <empty> ContainerId containerId; // launch an app to queue a1 (label = x), and check all container will // be allocated in h1 RMApp app1 = rm1.submitApp(200, "app", "user", null, "a1"); MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm3); // request a container. am1.allocate("*", 1024, 1, new ArrayList<ContainerId>(), "x"); containerId = ContainerId.newContainerId(am1.getApplicationAttemptId(), 2); Assert.assertFalse(rm1.waitForState(nm2, containerId, RMContainerState.ALLOCATED, 10 * 1000)); Assert.assertTrue(rm1.waitForState(nm1, containerId, RMContainerState.ALLOCATED, 10 * 1000)); checkTaskContainersHost(am1.getApplicationAttemptId(), containerId, rm1, "h1"); // launch an app to queue b1 (label = y), and check all container will // be allocated in h2 RMApp app2 = rm1.submitApp(200, "app", "user", null, "b1"); MockAM am2 = MockRM.launchAndRegisterAM(app2, rm1, nm3); // request a container. am2.allocate("*", 1024, 1, new ArrayList<ContainerId>(), "y"); containerId = ContainerId.newContainerId(am2.getApplicationAttemptId(), 2); Assert.assertFalse(rm1.waitForState(nm1, containerId, RMContainerState.ALLOCATED, 10 * 1000)); Assert.assertTrue(rm1.waitForState(nm2, containerId, RMContainerState.ALLOCATED, 10 * 1000)); checkTaskContainersHost(am2.getApplicationAttemptId(), containerId, rm1, "h2"); // launch an app to queue c1 (label = ""), and check all container will // be allocated in h3 RMApp app3 = rm1.submitApp(200, "app", "user", null, "c1"); MockAM am3 = MockRM.launchAndRegisterAM(app3, rm1, nm3); // request a container. am3.allocate("*", 1024, 1, new ArrayList<ContainerId>()); containerId = ContainerId.newContainerId(am3.getApplicationAttemptId(), 2); Assert.assertFalse(rm1.waitForState(nm2, containerId, RMContainerState.ALLOCATED, 10 * 1000)); Assert.assertTrue(rm1.waitForState(nm3, containerId, RMContainerState.ALLOCATED, 10 * 1000)); checkTaskContainersHost(am3.getApplicationAttemptId(), containerId, rm1, "h3"); rm1.close(); } @Test (timeout = 120000) public void testContainerAllocateWithDefaultQueueLabels() throws Exception { // This test is pretty much similar to testContainerAllocateWithLabel. // Difference is, this test doesn't specify label expression in ResourceRequest, // instead, it uses default queue label expression // set node -> label mgr.addToCluserNodeLabels(ImmutableSet.of("x", "y")); mgr.addLabelsToNode(ImmutableMap.of(NodeId.newInstance("h1", 0), toSet("x"), NodeId.newInstance("h2", 0), toSet("y"))); // inject node label manager MockRM rm1 = new MockRM(getConfigurationWithDefaultQueueLabels(conf)) { @Override public RMNodeLabelsManager createNodeLabelManager() { return mgr; } }; rm1.getRMContext().setNodeLabelManager(mgr); rm1.start(); MockNM nm1 = rm1.registerNode("h1:1234", 8000); // label = x MockNM nm2 = rm1.registerNode("h2:1234", 8000); // label = y MockNM nm3 = rm1.registerNode("h3:1234", 8000); // label = <empty> ContainerId containerId; // launch an app to queue a1 (label = x), and check all container will // be allocated in h1 RMApp app1 = rm1.submitApp(200, "app", "user", null, "a1"); MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1); // request a container. am1.allocate("*", 1024, 1, new ArrayList<ContainerId>()); containerId = ContainerId.newContainerId(am1.getApplicationAttemptId(), 2); Assert.assertFalse(rm1.waitForState(nm3, containerId, RMContainerState.ALLOCATED, 10 * 1000)); Assert.assertTrue(rm1.waitForState(nm1, containerId, RMContainerState.ALLOCATED, 10 * 1000)); checkTaskContainersHost(am1.getApplicationAttemptId(), containerId, rm1, "h1"); // launch an app to queue b1 (label = y), and check all container will // be allocated in h2 RMApp app2 = rm1.submitApp(200, "app", "user", null, "b1"); MockAM am2 = MockRM.launchAndRegisterAM(app2, rm1, nm2); // request a container. am2.allocate("*", 1024, 1, new ArrayList<ContainerId>()); containerId = ContainerId.newContainerId(am2.getApplicationAttemptId(), 2); Assert.assertFalse(rm1.waitForState(nm3, containerId, RMContainerState.ALLOCATED, 10 * 1000)); Assert.assertTrue(rm1.waitForState(nm2, containerId, RMContainerState.ALLOCATED, 10 * 1000)); checkTaskContainersHost(am2.getApplicationAttemptId(), containerId, rm1, "h2"); // launch an app to queue c1 (label = ""), and check all container will // be allocated in h3 RMApp app3 = rm1.submitApp(200, "app", "user", null, "c1"); MockAM am3 = MockRM.launchAndRegisterAM(app3, rm1, nm3); // request a container. am3.allocate("*", 1024, 1, new ArrayList<ContainerId>()); containerId = ContainerId.newContainerId(am3.getApplicationAttemptId(), 2); Assert.assertFalse(rm1.waitForState(nm2, containerId, RMContainerState.ALLOCATED, 10 * 1000)); Assert.assertTrue(rm1.waitForState(nm3, containerId, RMContainerState.ALLOCATED, 10 * 1000)); checkTaskContainersHost(am3.getApplicationAttemptId(), containerId, rm1, "h3"); rm1.close(); } }