Java MockRM类代码示例

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
private void verifyNewVersionToken(final Configuration conf, final CustomAM am,
    Token<ClientToAMTokenIdentifier> token, MockRM rm) throws IOException,
    InterruptedException {
  UserGroupInformation ugi;
  ugi = UserGroupInformation.createRemoteUser("me");
  
  Token<ClientToAMTokenIdentifier> newToken = 
      new Token<ClientToAMTokenIdentifier>(
          new ClientToAMTokenIdentifierForTest(token.decodeIdentifier(), "message"),
          am.getClientToAMTokenSecretManager());
  newToken.setService(token.getService());
  
  ugi.addToken(newToken);

  ugi.doAs(new PrivilegedExceptionAction<Void>() {
    @Override
    public Void run() throws Exception {
      CustomProtocol client =
          (CustomProtocol) RPC.getProxy(CustomProtocol.class, 1L, am.address,
            conf);
      client.ping();
      Assert.assertTrue(am.pinged);
      return null;
    }
  });
}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
@Override
protected void configureServlets() {
  bind(JAXBContextResolver.class);
  bind(RMWebServices.class);
  bind(GenericExceptionHandler.class);
  try {
    userName = UserGroupInformation.getCurrentUser().getShortUserName();
  } catch (IOException ioe) {
    throw new RuntimeException("Unable to get current user name "
        + ioe.getMessage(), ioe);
  }
  notUserName = userName + "abc123";
  conf = new YarnConfiguration();
  conf.set(YarnConfiguration.YARN_ADMIN_ACL, userName);
  rm = new MockRM(conf);
  bind(ResourceManager.class).toInstance(rm);
  filter("/*").through(
      TestRMWebServicesAppsModification.TestRMCustomAuthFilter.class);
  serve("/*").with(GuiceContainer.class);
}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
@Override
protected void configureServlets() {
  bind(JAXBContextResolver.class);
  bind(RMWebServices.class);
  bind(GenericExceptionHandler.class);
  Configuration rmconf = new Configuration();
  rmconf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS,
    YarnConfiguration.DEFAULT_RM_AM_MAX_ATTEMPTS);
  rmconf.setClass(YarnConfiguration.RM_SCHEDULER, FifoScheduler.class,
    ResourceScheduler.class);
  rmconf.setBoolean(YarnConfiguration.YARN_ACL_ENABLE, true);
  rm = new MockRM(rmconf);
  bind(ResourceManager.class).toInstance(rm);
  if (isKerberosAuth == true) {
    filter("/*").through(TestKerberosAuthFilter.class);
  } else {
    filter("/*").through(TestSimpleAuthFilter.class);
  }
  serve("/*").with(GuiceContainer.class);
}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
private static void setupAndStartRM() throws Exception {
  Configuration rmconf = new Configuration();
  rmconf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS,
    YarnConfiguration.DEFAULT_RM_AM_MAX_ATTEMPTS);
  rmconf.setClass(YarnConfiguration.RM_SCHEDULER, FifoScheduler.class,
    ResourceScheduler.class);
  rmconf.setBoolean(YarnConfiguration.YARN_ACL_ENABLE, true);
  rmconf.set(CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHENTICATION,
    "kerberos");
  rmconf.set("yarn.resourcemanager.principal", spnegoPrincipal);
  rmconf.set("yarn.resourcemanager.keytab",
      spnegoKeytabFile.getAbsolutePath());
  rmconf.setBoolean("mockrm.webapp.enabled", true);
  UserGroupInformation.setConfiguration(rmconf);
  rm = new MockRM(rmconf);
  rm.start();

}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
@Before
public void setup() {
  mockClock = new MockClock();
  conf = createConfiguration();
  resourceManager = new MockRM(conf);
  resourceManager.start();

  scheduler = (FairScheduler) resourceManager.getResourceScheduler();
  scheduler.setClock(mockClock);

  assertTrue(scheduler.isContinuousSchedulingEnabled());
  assertEquals(
      FairSchedulerConfiguration.DEFAULT_CONTINUOUS_SCHEDULING_SLEEP_MS,
      scheduler.getContinuousSchedulingSleepMs());
  assertEquals(mockClock, scheduler.getClock());
}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
@Test
public void testUpdateDemand() {
  conf.set(FairSchedulerConfiguration.ASSIGN_MULTIPLE, "false");
  resourceManager = new MockRM(conf);
  resourceManager.start();
  scheduler = (FairScheduler) resourceManager.getResourceScheduler();
  scheduler.allocConf = mock(AllocationConfiguration.class);

  String queueName = "root.queue1";
  when(scheduler.allocConf.getMaxResources(queueName)).thenReturn(maxResource);
  when(scheduler.allocConf.getMinResources(queueName)).thenReturn(Resources.none());
  FSLeafQueue schedulable = new FSLeafQueue(queueName, scheduler, null);

  FSAppAttempt app = mock(FSAppAttempt.class);
  Mockito.when(app.getDemand()).thenReturn(maxResource);

  schedulable.addAppSchedulable(app);
  schedulable.addAppSchedulable(app);

  schedulable.updateDemand();

  assertTrue("Demand is greater than max allowed ",
      Resources.equals(schedulable.getDemand(), maxResource));
}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
@Test
public void testAddAndRemoveAppFromCapacityScheduler() throws Exception {
  CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
  setupQueueConfiguration(conf);
  conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class,
    ResourceScheduler.class);
  MockRM rm = new MockRM(conf);
  @SuppressWarnings("unchecked")
  AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode> cs =
      (AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>) rm
        .getResourceScheduler();

  SchedulerApplication<SchedulerApplicationAttempt> app =
      TestSchedulerUtils.verifyAppAddedAndRemovedFromScheduler(
        cs.getSchedulerApplications(), cs, "a1");
  Assert.assertEquals("a1", app.getQueue().getQueueName());
}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
@Test
public void testAsyncScheduling() throws Exception {
  Configuration conf = new Configuration();
  conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class,
      ResourceScheduler.class);
  MockRM rm = new MockRM(conf);
  rm.start();
  CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();

  final int NODES = 100;
  
  // Register nodes
  for (int i=0; i < NODES; ++i) {
    String host = "192.168.1." + i;
    RMNode node =
        MockNodes.newNodeInfo(0, MockNodes.newResource(4 * GB), 1, host);
    cs.handle(new NodeAddedSchedulerEvent(node));
  }
  
  // Now directly exercise the scheduling loop
  for (int i=0; i < NODES; ++i) {
    CapacityScheduler.schedule(cs);
  }
}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
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();
}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
private void checkAppQueue(MockRM resourceManager, String user,
    String submissionQueue, String expected)
    throws Exception {
  RMApp app = resourceManager.submitApp(200, "name", user,
      new HashMap<ApplicationAccessType, String>(), false, submissionQueue, -1,
      null, "MAPREDUCE", false);
  RMAppState expectedState = expected.isEmpty() ? RMAppState.FAILED
      : RMAppState.ACCEPTED;
  resourceManager.waitForState(app.getApplicationId(), expectedState);
  // get scheduler app
  CapacityScheduler cs = (CapacityScheduler)
      resourceManager.getResourceScheduler();
  SchedulerApplication schedulerApp =
      cs.getSchedulerApplications().get(app.getApplicationId());
  String queue = "";
  if (schedulerApp != null) {
    queue = schedulerApp.getQueue().getQueueName();
  }
  Assert.assertTrue("expected " + expected + " actual " + queue,
      expected.equals(queue));
  Assert.assertEquals(expected, app.getQueue());
}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
@Test
public void testPolicyInitializeAfterSchedulerInitialized() {
  @SuppressWarnings("resource")
  MockRM rm = new MockRM(conf);
  rm.init(conf);
  
  // ProportionalCapacityPreemptionPolicy should be initialized after
  // CapacityScheduler initialized. We will 
  // 1) find SchedulingMonitor from RMActiveService's service list, 
  // 2) check if ResourceCalculator in policy is null or not. 
  // If it's not null, we can come to a conclusion that policy initialized
  // after scheduler got initialized
  for (Service service : rm.getRMActiveService().getServices()) {
    if (service instanceof SchedulingMonitor) {
      ProportionalCapacityPreemptionPolicy policy =
          (ProportionalCapacityPreemptionPolicy) ((SchedulingMonitor) service)
              .getSchedulingEditPolicy();
      assertNotNull(policy.getResourceCalculator());
      return;
    }
  }
  
  fail("Failed to find SchedulingMonitor service, please check what happened");
}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
@Test
public void testKillApplication() throws Exception {
  MockRM rm = new MockRM();
  rm.start();
  RMApp app = rm.submitApp(2000);

  Configuration conf = new Configuration();
  @SuppressWarnings("resource")
  final YarnClient client = new MockYarnClient();
  client.init(conf);
  client.start();

  client.killApplication(app.getApplicationId());
  verify(((MockYarnClient) client).getRMClient(), times(2))
    .forceKillApplication(any(KillApplicationRequest.class));
}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
@Before
public void setUp() throws IOException {
  scheduler = new FairScheduler();
  
  Configuration conf = new YarnConfiguration();
  conf.setClass(YarnConfiguration.RM_SCHEDULER, FairScheduler.class,
      ResourceScheduler.class);
  conf.set("yarn.scheduler.fair.event-log-enabled", "true");

  // All tests assume only one assignment per node update
  conf.set(FairSchedulerConfiguration.ASSIGN_MULTIPLE, "false");
  resourceManager = new MockRM(conf);
  ((AsyncDispatcher)resourceManager.getRMContext().getDispatcher()).start();
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, resourceManager.getRMContext());
}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
@Before
public void setup() {
  mockClock = new ControlledClock();
  conf = createConfiguration();
  resourceManager = new MockRM(conf);
  resourceManager.start();

  scheduler = (FairScheduler) resourceManager.getResourceScheduler();
  scheduler.setClock(mockClock);

  assertTrue(scheduler.isContinuousSchedulingEnabled());
  assertEquals(
      FairSchedulerConfiguration.DEFAULT_CONTINUOUS_SCHEDULING_SLEEP_MS,
      scheduler.getContinuousSchedulingSleepMs());
  assertEquals(mockClock, scheduler.getClock());
}
 

import org.apache.hadoop.yarn.server.resourcemanager.MockRM; //导入依赖的package包/类
private void allocateAndLaunchContainers(MockAM am, MockNM nm, MockRM rm,
    int nContainer, int mem, int priority, int startContainerId)
        throws Exception {
  am.allocate(Arrays
      .asList(ResourceRequest.newInstance(Priority.newInstance(priority), "*",
          Resources.createResource(mem), nContainer)),
      null);
  ContainerId lastContainerId = ContainerId.newContainerId(
      am.getApplicationAttemptId(), startContainerId + nContainer - 1);
  Assert.assertTrue(rm.waitForState(nm, lastContainerId,
      RMContainerState.ALLOCATED, 10 * 1000));
  // Acquire them, and NM report RUNNING
  am.allocate(null, null);

  for (int cId = startContainerId; cId < startContainerId
      + nContainer; cId++) {
    sentRMContainerLaunched(rm,
        ContainerId.newContainerId(am.getApplicationAttemptId(), cId));
    rm.waitForContainerState(
        ContainerId.newContainerId(am.getApplicationAttemptId(), cId),
        RMContainerState.RUNNING);
  }
}