org.apache.hadoop.yarn.api.ContainerManagementProtocol Java Examples

private void startContainer(final YarnRPC rpc,
    org.apache.hadoop.yarn.api.records.Token nmToken,
    org.apache.hadoop.yarn.api.records.Token containerToken,
    NodeId nodeId, String user) throws Exception {

  ContainerLaunchContext context =
      Records.newRecord(ContainerLaunchContext.class);
  StartContainerRequest scRequest =
      StartContainerRequest.newInstance(context,containerToken);
  List<StartContainerRequest> list = new ArrayList<StartContainerRequest>();
  list.add(scRequest);
  StartContainersRequest allRequests =
      StartContainersRequest.newInstance(list);
  ContainerManagementProtocol proxy = null;
  try {
    proxy = getContainerManagementProtocolProxy(rpc, nmToken, nodeId, user);
    StartContainersResponse response = proxy.startContainers(allRequests);
    for(SerializedException ex : response.getFailedRequests().values()){
      parseAndThrowException(ex.deSerialize());
    }
  } finally {
    if (proxy != null) {
      rpc.stopProxy(proxy, conf);
    }
  }
}
 

private void stopContainer(YarnRPC rpc, Token nmToken,
    List<ContainerId> containerId, ApplicationAttemptId appAttemptId,
    NodeId nodeId) throws Exception {
  StopContainersRequest request =
      StopContainersRequest.newInstance(containerId);
  ContainerManagementProtocol proxy = null;
  try {
    proxy =
        getContainerManagementProtocolProxy(rpc, nmToken, nodeId,
            appAttemptId.toString());
    StopContainersResponse response = proxy.stopContainers(request);
    if (response.getFailedRequests() != null &&
        response.getFailedRequests().containsKey(containerId)) {
      parseAndThrowException(response.getFailedRequests().get(containerId)
          .deSerialize());
    }
  } catch (Exception e) {
    if (proxy != null) {
      rpc.stopProxy(proxy, conf);
    }
  }
}
 

private void startContainer(final YarnRPC rpc,
    org.apache.hadoop.yarn.api.records.Token nmToken,
    org.apache.hadoop.yarn.api.records.Token containerToken,
    NodeId nodeId, String user) throws Exception {

  ContainerLaunchContext context =
      Records.newRecord(ContainerLaunchContext.class);
  StartContainerRequest scRequest =
      StartContainerRequest.newInstance(context,containerToken);
  List<StartContainerRequest> list = new ArrayList<StartContainerRequest>();
  list.add(scRequest);
  StartContainersRequest allRequests =
      StartContainersRequest.newInstance(list);
  ContainerManagementProtocol proxy = null;
  try {
    proxy = getContainerManagementProtocolProxy(rpc, nmToken, nodeId, user);
    StartContainersResponse response = proxy.startContainers(allRequests);
    for(SerializedException ex : response.getFailedRequests().values()){
      parseAndThrowException(ex.deSerialize());
    }
  } finally {
    if (proxy != null) {
      rpc.stopProxy(proxy, conf);
    }
  }
}
 

private void stopContainer(YarnRPC rpc, Token nmToken,
    List<ContainerId> containerId, ApplicationAttemptId appAttemptId,
    NodeId nodeId) throws Exception {
  StopContainersRequest request =
      StopContainersRequest.newInstance(containerId);
  ContainerManagementProtocol proxy = null;
  try {
    proxy =
        getContainerManagementProtocolProxy(rpc, nmToken, nodeId,
            appAttemptId.toString());
    StopContainersResponse response = proxy.stopContainers(request);
    if (response.getFailedRequests() != null &&
        response.getFailedRequests().containsKey(containerId)) {
      parseAndThrowException(response.getFailedRequests().get(containerId)
          .deSerialize());
    }
  } catch (Exception e) {
    if (proxy != null) {
      rpc.stopProxy(proxy, conf);
    }
  }
}
 

public static void waitForContainerState(ContainerManagementProtocol containerManager,
      ContainerId containerID, ContainerState finalState, int timeOutMax)
      throws InterruptedException, YarnException, IOException {
List<ContainerId> list = new ArrayList<ContainerId>();
list.add(containerID);
GetContainerStatusesRequest request =
    GetContainerStatusesRequest.newInstance(list);
ContainerStatus containerStatus =
    containerManager.getContainerStatuses(request).getContainerStatuses()
      .get(0);
int timeoutSecs = 0;
  while (!containerStatus.getState().equals(finalState)
      && timeoutSecs++ < timeOutMax) {
      Thread.sleep(1000);
      LOG.info("Waiting for container to get into state " + finalState
          + ". Current state is " + containerStatus.getState());
      containerStatus = containerManager.getContainerStatuses(request).getContainerStatuses().get(0);
    }
    LOG.info("Container state is " + containerStatus.getState());
    Assert.assertEquals("ContainerState is not correct (timedout)",
        finalState, containerStatus.getState());
  }
 

@Private
@VisibleForTesting
protected ContainerManagementProtocol newProxy(final YarnRPC rpc,
    String containerManagerBindAddr, ContainerId containerId, Token token)
    throws InvalidToken {

  if (token == null) {
    throw new InvalidToken("No NMToken sent for "
        + containerManagerBindAddr);
  }
  
  final InetSocketAddress cmAddr =
      NetUtils.createSocketAddr(containerManagerBindAddr);
  LOG.info("Opening proxy : " + containerManagerBindAddr);
  // the user in createRemoteUser in this context has to be ContainerID
  UserGroupInformation user =
      UserGroupInformation.createRemoteUser(containerId
          .getApplicationAttemptId().toString());

  org.apache.hadoop.security.token.Token<NMTokenIdentifier> nmToken =
      ConverterUtils.convertFromYarn(token, cmAddr);
  user.addToken(nmToken);

  return NMProxy.createNMProxy(conf, ContainerManagementProtocol.class,
    user, rpc, cmAddr);
}
 

@Private
@VisibleForTesting
protected ContainerManagementProtocol newProxy(final YarnRPC rpc,
    String containerManagerBindAddr, ContainerId containerId, Token token)
    throws InvalidToken {

  if (token == null) {
    throw new InvalidToken("No NMToken sent for "
        + containerManagerBindAddr);
  }
  
  final InetSocketAddress cmAddr =
      NetUtils.createSocketAddr(containerManagerBindAddr);
  LOG.info("Opening proxy : " + containerManagerBindAddr);
  // the user in createRemoteUser in this context has to be ContainerID
  UserGroupInformation user =
      UserGroupInformation.createRemoteUser(containerId
          .getApplicationAttemptId().toString());

  org.apache.hadoop.security.token.Token<NMTokenIdentifier> nmToken =
      ConverterUtils.convertFromYarn(token, cmAddr);
  user.addToken(nmToken);

  return NMProxy.createNMProxy(conf, ContainerManagementProtocol.class,
    user, rpc, cmAddr);
}
 

public static void waitForContainerState(ContainerManagementProtocol containerManager,
      ContainerId containerID, ContainerState finalState, int timeOutMax)
      throws InterruptedException, YarnException, IOException {
List<ContainerId> list = new ArrayList<ContainerId>();
list.add(containerID);
GetContainerStatusesRequest request =
    GetContainerStatusesRequest.newInstance(list);
ContainerStatus containerStatus =
    containerManager.getContainerStatuses(request).getContainerStatuses()
      .get(0);
int timeoutSecs = 0;
  while (!containerStatus.getState().equals(finalState)
      && timeoutSecs++ < timeOutMax) {
      Thread.sleep(1000);
      LOG.info("Waiting for container to get into state " + finalState
          + ". Current state is " + containerStatus.getState());
      containerStatus = containerManager.getContainerStatuses(request).getContainerStatuses().get(0);
    }
    LOG.info("Container state is " + containerStatus.getState());
    Assert.assertEquals("ContainerState is not correct (timedout)",
        finalState, containerStatus.getState());
  }
 

protected ContainerManagementProtocol getContainerMgrProxy(
    final ContainerId containerId) {

  final NodeId node = masterContainer.getNodeId();
  final InetSocketAddress containerManagerBindAddress =
      NetUtils.createSocketAddrForHost(node.getHost(), node.getPort());

  final YarnRPC rpc = YarnRPC.create(conf); // TODO: Don't create again and again.

  UserGroupInformation currentUser =
      UserGroupInformation.createRemoteUser(containerId
          .getApplicationAttemptId().toString());

  String user =
      rmContext.getRMApps()
          .get(containerId.getApplicationAttemptId().getApplicationId())
          .getUser();
  org.apache.hadoop.yarn.api.records.Token token =
      rmContext.getNMTokenSecretManager().createNMToken(
          containerId.getApplicationAttemptId(), node, user);
  currentUser.addToken(ConverterUtils.convertFromYarn(token,
      containerManagerBindAddress));

  return currentUser
      .doAs(new PrivilegedAction<ContainerManagementProtocol>() {

        @Override
        public ContainerManagementProtocol run() {
          return (ContainerManagementProtocol) rpc.getProxy(
              ContainerManagementProtocol.class,
              containerManagerBindAddress, conf);
        }
      });
}
 

protected ContainerManagementProtocol getContainerManagementProtocolProxy(
    final YarnRPC rpc, org.apache.hadoop.yarn.api.records.Token nmToken,
    NodeId nodeId, String user) {
  ContainerManagementProtocol proxy;
  UserGroupInformation ugi = UserGroupInformation.createRemoteUser(user);
  final InetSocketAddress addr =
      NetUtils.createSocketAddr(nodeId.getHost(), nodeId.getPort());
  if (nmToken != null) {
    ugi.addToken(ConverterUtils.convertFromYarn(nmToken, addr));      
  }
  proxy =
      NMProxy.createNMProxy(conf, ContainerManagementProtocol.class, ugi,
        rpc, addr);
  return proxy;
}
 

private void
    getContainerStatus(YarnRPC rpc,
        org.apache.hadoop.yarn.api.records.Token nmToken,
        ContainerId containerId,
        ApplicationAttemptId appAttemptId, NodeId nodeId,
        boolean isExceptionExpected) throws Exception {
  List<ContainerId> containerIds = new ArrayList<ContainerId>();
  containerIds.add(containerId);
  GetContainerStatusesRequest request =
      GetContainerStatusesRequest.newInstance(containerIds);
  ContainerManagementProtocol proxy = null;
  try {
    proxy =
        getContainerManagementProtocolProxy(rpc, nmToken, nodeId,
            appAttemptId.toString());
    GetContainerStatusesResponse statuses = proxy.getContainerStatuses(request);
    if (statuses.getFailedRequests() != null
        && statuses.getFailedRequests().containsKey(containerId)) {
      parseAndThrowException(statuses.getFailedRequests().get(containerId)
        .deSerialize());
    }
  } finally {
    if (proxy != null) {
      rpc.stopProxy(proxy, conf);
    }
  }
}
 

@Override
protected ApplicationMasterLauncher createAMLauncher() {
  return new ApplicationMasterLauncher(getRMContext()) {
    @Override
    protected Runnable createRunnableLauncher(RMAppAttempt application,
        AMLauncherEventType event) {
      return new AMLauncher(context, application, event, getConfig()) {
        @Override
        protected ContainerManagementProtocol getContainerMgrProxy(
            ContainerId containerId) {
          return containerManager;
        }
        @Override
        protected Token<AMRMTokenIdentifier> createAndSetAMRMToken() {
          Token<AMRMTokenIdentifier> amRmToken =
              super.createAndSetAMRMToken();
          InetSocketAddress serviceAddr =
              getConfig().getSocketAddr(
                YarnConfiguration.RM_SCHEDULER_ADDRESS,
                YarnConfiguration.DEFAULT_RM_SCHEDULER_ADDRESS,
                YarnConfiguration.DEFAULT_RM_SCHEDULER_PORT);
          SecurityUtil.setTokenService(amRmToken, serviceAddr);
          return amRmToken;
        }
      };
    }
  };
}
 

protected ContainerManagementProtocol getContainerMgrProxy(
    final ContainerId containerId) {

  final NodeId node = masterContainer.getNodeId();
  final InetSocketAddress containerManagerBindAddress =
      NetUtils.createSocketAddrForHost(node.getHost(), node.getPort());

  final YarnRPC rpc = YarnRPC.create(conf); // TODO: Don't create again and again.

  UserGroupInformation currentUser =
      UserGroupInformation.createRemoteUser(containerId
          .getApplicationAttemptId().toString());

  String user =
      rmContext.getRMApps()
          .get(containerId.getApplicationAttemptId().getApplicationId())
          .getUser();
  org.apache.hadoop.yarn.api.records.Token token =
      rmContext.getNMTokenSecretManager().createNMToken(
          containerId.getApplicationAttemptId(), node, user);
  currentUser.addToken(ConverterUtils.convertFromYarn(token,
      containerManagerBindAddress));

  return currentUser
      .doAs(new PrivilegedAction<ContainerManagementProtocol>() {

        @Override
        public ContainerManagementProtocol run() {
          return (ContainerManagementProtocol) rpc.getProxy(
              ContainerManagementProtocol.class,
              containerManagerBindAddress, conf);
        }
      });
}
 

@Override
protected ApplicationMasterLauncher createAMLauncher() {
  return new ApplicationMasterLauncher(getRMContext()) {
    @Override
    protected Runnable createRunnableLauncher(RMAppAttempt application,
        AMLauncherEventType event) {
      return new AMLauncher(context, application, event, getConfig()) {
        @Override
        protected ContainerManagementProtocol getContainerMgrProxy(
            ContainerId containerId) {
          return containerManager;
        }
        @Override
        protected Token<AMRMTokenIdentifier> createAndSetAMRMToken() {
          Token<AMRMTokenIdentifier> amRmToken =
              super.createAndSetAMRMToken();
          InetSocketAddress serviceAddr =
              getConfig().getSocketAddr(
                YarnConfiguration.RM_SCHEDULER_ADDRESS,
                YarnConfiguration.DEFAULT_RM_SCHEDULER_ADDRESS,
                YarnConfiguration.DEFAULT_RM_SCHEDULER_PORT);
          SecurityUtil.setTokenService(amRmToken, serviceAddr);
          return amRmToken;
        }
      };
    }
  };
}
 

private void
    getContainerStatus(YarnRPC rpc,
        org.apache.hadoop.yarn.api.records.Token nmToken,
        ContainerId containerId,
        ApplicationAttemptId appAttemptId, NodeId nodeId,
        boolean isExceptionExpected) throws Exception {
  List<ContainerId> containerIds = new ArrayList<ContainerId>();
  containerIds.add(containerId);
  GetContainerStatusesRequest request =
      GetContainerStatusesRequest.newInstance(containerIds);
  ContainerManagementProtocol proxy = null;
  try {
    proxy =
        getContainerManagementProtocolProxy(rpc, nmToken, nodeId,
            appAttemptId.toString());
    GetContainerStatusesResponse statuses = proxy.getContainerStatuses(request);
    if (statuses.getFailedRequests() != null
        && statuses.getFailedRequests().containsKey(containerId)) {
      parseAndThrowException(statuses.getFailedRequests().get(containerId)
        .deSerialize());
    }
  } finally {
    if (proxy != null) {
      rpc.stopProxy(proxy, conf);
    }
  }
}
 

protected ContainerManagementProtocol getContainerManagementProtocolProxy(
    final YarnRPC rpc, org.apache.hadoop.yarn.api.records.Token nmToken,
    NodeId nodeId, String user) {
  ContainerManagementProtocol proxy;
  UserGroupInformation ugi = UserGroupInformation.createRemoteUser(user);
  final InetSocketAddress addr =
      NetUtils.createSocketAddr(nodeId.getHost(), nodeId.getPort());
  if (nmToken != null) {
    ugi.addToken(ConverterUtils.convertFromYarn(nmToken, addr));      
  }
  proxy =
      NMProxy.createNMProxy(conf, ContainerManagementProtocol.class, ugi,
        rpc, addr);
  return proxy;
}
 

protected ContainerManagementProtocol getCMProxy(ContainerId containerID,
                                                 final String containerManagerBindAddr,
                                                 Token containerToken)
    throws IOException {
  String [] hosts = containerManagerBindAddr.split(":");
  final InetSocketAddress cmAddr =
      new InetSocketAddress(hosts[0], Integer.parseInt(hosts[1]));
  UserGroupInformation user = UserGroupInformation.getCurrentUser();

  if (UserGroupInformation.isSecurityEnabled()) {
    org.apache.hadoop.security.token.Token<ContainerTokenIdentifier> token =
        ConverterUtils.convertFromYarn(containerToken, cmAddr);
    // the user in createRemoteUser in this context has to be ContainerID
    user = UserGroupInformation.createRemoteUser(containerID.toString());
    user.addToken(token);
  }

  ContainerManagementProtocol proxy = user.doAs(new PrivilegedAction<ContainerManagementProtocol>() {
    @Override
    public ContainerManagementProtocol run() {
      return (ContainerManagementProtocol) yarnRPC.getProxy(ContainerManagementProtocol.class,
          cmAddr, conf);
    }
  });

  return proxy;
}
 

@Override
  public synchronized void stopContainer() {

    if(isCompletelyDone()) {
      return;
    }
    if(this.state == ContainerState.PREP) {
      this.state = ContainerState.KILLED_BEFORE_LAUNCH;
    } else {
      LOG.info("KILLING " + containerID);

      ContainerManagementProtocol proxy = null;
      try {
        proxy = getCMProxy(this.containerID, this.containerMgrAddress,
            this.containerToken);

        // kill the remote container if already launched
        List<ContainerId> willBeStopedIds = new ArrayList<ContainerId>();
        willBeStopedIds.add(this.containerID);
        StopContainersRequest stopRequests = Records.newRecord(StopContainersRequest.class);
        stopRequests.setContainerIds(willBeStopedIds);
        proxy.stopContainers(stopRequests);
        // If stopContainer returns without an error, assuming the stop made
        // it over to the NodeManager.
//          context.getEventHandler().handle(
//              new AMContainerEvent(containerID, AMContainerEventType.C_NM_STOP_SENT));
        context.getResourceAllocator().removeContainer(containerID);
      } catch (Throwable t) {

        // ignore the cleanup failure
        String message = "cleanup failed for container "
            + this.containerID + " : "
            + StringUtils.stringifyException(t);
//          context.getEventHandler().handle(
//              new AMContainerEventStopFailed(containerID, message));
        LOG.warn(message);
        this.state = ContainerState.DONE;
        return;
      } finally {
        if (proxy != null) {
          yarnRPC.stopProxy(proxy, conf);
        }
      }
      this.state = ContainerState.DONE;
    }
  }
 

@Override
@SuppressWarnings("unchecked")
public synchronized void launch(ContainerLaunchContext commonContainerLaunchContext) {
  LOG.info("Launching Container with Id: " + containerID);
  if(this.state == ContainerState.KILLED_BEFORE_LAUNCH) {
    state = ContainerState.DONE;
    LOG.error("Container (" + containerID + " was killed before it was launched");
    return;
  }

  ContainerManagementProtocol proxy = null;
  try {

    proxy = getCMProxy(containerID, containerMgrAddress,
        containerToken);

    // Construct the actual Container
    ContainerLaunchContext containerLaunchContext = createContainerLaunchContext(commonContainerLaunchContext);

    // Now launch the actual container
    List<StartContainerRequest> startRequestList = new ArrayList<StartContainerRequest>();
    StartContainerRequest startRequest = Records
        .newRecord(StartContainerRequest.class);
    startRequest.setContainerLaunchContext(containerLaunchContext);
    startRequestList.add(startRequest);
    StartContainersRequest startRequests = Records.newRecord(StartContainersRequest.class);
    startRequests.setStartContainerRequests(startRequestList);
    StartContainersResponse response = proxy.startContainers(startRequests);

    ByteBuffer portInfo = response.getAllServicesMetaData().get(PullServerAuxService.PULLSERVER_SERVICEID);

    if(portInfo != null) {
      port = PullServerAuxService.deserializeMetaData(portInfo);
    }

    LOG.info("PullServer port returned by ContainerManager for "
        + containerID + " : " + port);

    if(port < 0) {
      this.state = ContainerState.FAILED;
      throw new IllegalStateException("Invalid shuffle port number "
          + port + " returned for " + containerID);
    }

    this.state = ContainerState.RUNNING;
    this.hostName = containerMgrAddress.split(":")[0];
    context.getResourceAllocator().addContainer(containerID, this);
  } catch (Throwable t) {
    String message = "Container launch failed for " + containerID + " : "
        + StringUtils.stringifyException(t);
    this.state = ContainerState.FAILED;
    LOG.error(message);
  } finally {
    if (proxy != null) {
      yarnRPC.stopProxy(proxy, conf);
    }
  }
}
 

@Override
protected void serviceStart() throws Exception {

  // Enqueue user dirs in deletion context

  Configuration conf = getConfig();
  final InetSocketAddress initialAddress = conf.getSocketAddr(
      YarnConfiguration.NM_BIND_HOST,
      YarnConfiguration.NM_ADDRESS,
      YarnConfiguration.DEFAULT_NM_ADDRESS,
      YarnConfiguration.DEFAULT_NM_PORT);
  boolean usingEphemeralPort = (initialAddress.getPort() == 0);
  if (context.getNMStateStore().canRecover() && usingEphemeralPort) {
    throw new IllegalArgumentException("Cannot support recovery with an "
        + "ephemeral server port. Check the setting of "
        + YarnConfiguration.NM_ADDRESS);
  }
  // If recovering then delay opening the RPC service until the recovery
  // of resources and containers have completed, otherwise requests from
  // clients during recovery can interfere with the recovery process.
  final boolean delayedRpcServerStart =
      context.getNMStateStore().canRecover();

  Configuration serverConf = new Configuration(conf);

  // always enforce it to be token-based.
  serverConf.set(
    CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHENTICATION,
    SaslRpcServer.AuthMethod.TOKEN.toString());
  
  YarnRPC rpc = YarnRPC.create(conf);

  server =
      rpc.getServer(ContainerManagementProtocol.class, this, initialAddress, 
          serverConf, this.context.getNMTokenSecretManager(),
          conf.getInt(YarnConfiguration.NM_CONTAINER_MGR_THREAD_COUNT, 
              YarnConfiguration.DEFAULT_NM_CONTAINER_MGR_THREAD_COUNT));
  
  // Enable service authorization?
  if (conf.getBoolean(
      CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHORIZATION, 
      false)) {
    refreshServiceAcls(conf, new NMPolicyProvider());
  }
  
  LOG.info("Blocking new container-requests as container manager rpc" +
  		" server is still starting.");
  this.setBlockNewContainerRequests(true);

  String bindHost = conf.get(YarnConfiguration.NM_BIND_HOST);
  String nmAddress = conf.getTrimmed(YarnConfiguration.NM_ADDRESS);
  String hostOverride = null;
  if (bindHost != null && !bindHost.isEmpty()
      && nmAddress != null && !nmAddress.isEmpty()) {
    //a bind-host case with an address, to support overriding the first
    //hostname found when querying for our hostname with the specified
    //address, combine the specified address with the actual port listened
    //on by the server
    hostOverride = nmAddress.split(":")[0];
  }

  // setup node ID
  InetSocketAddress connectAddress;
  if (delayedRpcServerStart) {
    connectAddress = NetUtils.getConnectAddress(initialAddress);
  } else {
    server.start();
    connectAddress = NetUtils.getConnectAddress(server);
  }
  NodeId nodeId = buildNodeId(connectAddress, hostOverride);
  ((NodeManager.NMContext)context).setNodeId(nodeId);
  this.context.getNMTokenSecretManager().setNodeId(nodeId);
  this.context.getContainerTokenSecretManager().setNodeId(nodeId);

  // start remaining services
  super.serviceStart();

  if (delayedRpcServerStart) {
    waitForRecoveredContainers();
    server.start();

    // check that the node ID is as previously advertised
    connectAddress = NetUtils.getConnectAddress(server);
    NodeId serverNode = buildNodeId(connectAddress, hostOverride);
    if (!serverNode.equals(nodeId)) {
      throw new IOException("Node mismatch after server started, expected '"
          + nodeId + "' but found '" + serverNode + "'");
    }
  }

  LOG.info("ContainerManager started at " + connectAddress);
  LOG.info("ContainerManager bound to " + initialAddress);
}
 

@Override
public ContainerManagementProtocol getContainerManager() {
  return this.containerManager;
}
 

public void setContainerManager(ContainerManagementProtocol containerManager) {
  this.containerManager = containerManager;
}
 

public static void waitForContainerState(ContainerManagementProtocol containerManager,
    ContainerId containerID, ContainerState finalState)
    throws InterruptedException, YarnException, IOException {
  waitForContainerState(containerManager, containerID, finalState, 20);
}
 

@Test(timeout = 20000)
public void testNMProxyRetry() throws Exception {
  containerManager.start();
  containerManager.setBlockNewContainerRequests(false);
  StartContainersRequest allRequests =
      Records.newRecord(StartContainersRequest.class);
  ApplicationId appId = ApplicationId.newInstance(1, 1);
  ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(appId, 1);

  org.apache.hadoop.yarn.api.records.Token nmToken =
      context.getNMTokenSecretManager().createNMToken(attemptId,
        context.getNodeId(), user);
  final InetSocketAddress address =
      conf.getSocketAddr(YarnConfiguration.NM_BIND_HOST,
        YarnConfiguration.NM_ADDRESS, YarnConfiguration.DEFAULT_NM_ADDRESS,
        YarnConfiguration.DEFAULT_NM_PORT);
  Token<NMTokenIdentifier> token =
      ConverterUtils.convertFromYarn(nmToken,
        SecurityUtil.buildTokenService(address));
  UserGroupInformation ugi = UserGroupInformation.createRemoteUser(user);
  ugi.addToken(token);

  ContainerManagementProtocol proxy =
      NMProxy.createNMProxy(conf, ContainerManagementProtocol.class, ugi,
        YarnRPC.create(conf), address);

  retryCount = 0;
  shouldThrowNMNotYetReadyException = false;
  proxy.startContainers(allRequests);
  Assert.assertEquals(5, retryCount);

  retryCount = 0;
  shouldThrowNMNotYetReadyException = false;
  proxy.stopContainers(Records.newRecord(StopContainersRequest.class));
  Assert.assertEquals(5, retryCount);

  retryCount = 0;
  shouldThrowNMNotYetReadyException = false;
  proxy.getContainerStatuses(Records
    .newRecord(GetContainerStatusesRequest.class));
  Assert.assertEquals(5, retryCount);

  retryCount = 0;
  shouldThrowNMNotYetReadyException = true;
  proxy.startContainers(allRequests);
  Assert.assertEquals(5, retryCount);
}
 

public ContainerLauncherImplUnderTest(AppContext context,
    ContainerManagementProtocol containerManager) {
  super(context);
  this.containerManager = containerManager;
}
 

@Test(timeout = 15000)
public void testSlowNM() throws Exception {

  conf = new Configuration();
  int maxAttempts = 1;
  conf.setInt(MRJobConfig.MAP_MAX_ATTEMPTS, maxAttempts);
  conf.setBoolean(MRJobConfig.JOB_UBERTASK_ENABLE, false);
  // set timeout low for the test
  conf.setInt("yarn.rpc.nm-command-timeout", 3000);
  conf.set(YarnConfiguration.IPC_RPC_IMPL, HadoopYarnProtoRPC.class.getName());
  YarnRPC rpc = YarnRPC.create(conf);
  String bindAddr = "localhost:0";
  InetSocketAddress addr = NetUtils.createSocketAddr(bindAddr);
  NMTokenSecretManagerInNM tokenSecretManager =
      new NMTokenSecretManagerInNM();
  MasterKey masterKey = Records.newRecord(MasterKey.class);
  masterKey.setBytes(ByteBuffer.wrap("key".getBytes()));
  tokenSecretManager.setMasterKey(masterKey);
  conf.set(CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHENTICATION,
    "token");
  server =
      rpc.getServer(ContainerManagementProtocol.class,
        new DummyContainerManager(), addr, conf, tokenSecretManager, 1);
  server.start();

  MRApp app = new MRAppWithSlowNM(tokenSecretManager);

  try {
  Job job = app.submit(conf);
  app.waitForState(job, JobState.RUNNING);

  Map<TaskId, Task> tasks = job.getTasks();
  Assert.assertEquals("Num tasks is not correct", 1, tasks.size());

  Task task = tasks.values().iterator().next();
  app.waitForState(task, TaskState.SCHEDULED);

  Map<TaskAttemptId, TaskAttempt> attempts = tasks.values().iterator()
      .next().getAttempts();
    Assert.assertEquals("Num attempts is not correct", maxAttempts,
        attempts.size());

  TaskAttempt attempt = attempts.values().iterator().next();
    app.waitForInternalState((TaskAttemptImpl) attempt,
        TaskAttemptStateInternal.ASSIGNED);

  app.waitForState(job, JobState.FAILED);

  String diagnostics = attempt.getDiagnostics().toString();
  LOG.info("attempt.getDiagnostics: " + diagnostics);

    Assert.assertTrue(diagnostics.contains("Container launch failed for "
        + "container_0_0000_01_000000 : "));
    Assert
        .assertTrue(diagnostics
            .contains("java.net.SocketTimeoutException: 3000 millis timeout while waiting for channel"));

  } finally {
    server.stop();
  app.stop();
}
}
 

public ContainerManagementProtocol getContainerManagementProtocol() {
  return proxy;
}
 

public MockRMWithAMS(Configuration conf, ContainerManagementProtocol containerManager) {
  super(conf, containerManager);
}
 

public MockRMWithCustomAMLauncher(ContainerManagementProtocol containerManager) {
  this(new Configuration(), containerManager);
}
 

private void test(String rpcClass) throws Exception {
  Configuration conf = new Configuration();
  conf.set(YarnConfiguration.IPC_RPC_IMPL, rpcClass);
  YarnRPC rpc = YarnRPC.create(conf);
  String bindAddr = "localhost:0";
  InetSocketAddress addr = NetUtils.createSocketAddr(bindAddr);
  Server server = rpc.getServer(ContainerManagementProtocol.class, 
          new DummyContainerManager(), addr, conf, null, 1);
  server.start();
  RPC.setProtocolEngine(conf, ContainerManagementProtocolPB.class, ProtobufRpcEngine.class);
  ContainerManagementProtocol proxy = (ContainerManagementProtocol) 
      rpc.getProxy(ContainerManagementProtocol.class, 
          NetUtils.getConnectAddress(server), conf);
  ContainerLaunchContext containerLaunchContext = 
      recordFactory.newRecordInstance(ContainerLaunchContext.class);

  ApplicationId applicationId = ApplicationId.newInstance(0, 0);
  ApplicationAttemptId applicationAttemptId =
      ApplicationAttemptId.newInstance(applicationId, 0);
  ContainerId containerId =
      ContainerId.newContainerId(applicationAttemptId, 100);
  NodeId nodeId = NodeId.newInstance("localhost", 1234);
  Resource resource = Resource.newInstance(1234, 2);
  ContainerTokenIdentifier containerTokenIdentifier =
      new ContainerTokenIdentifier(containerId, "localhost", "user",
        resource, System.currentTimeMillis() + 10000, 42, 42,
        Priority.newInstance(0), 0);
  Token containerToken = newContainerToken(nodeId, "password".getBytes(),
        containerTokenIdentifier);

  StartContainerRequest scRequest =
      StartContainerRequest.newInstance(containerLaunchContext,
        containerToken);
  List<StartContainerRequest> list = new ArrayList<StartContainerRequest>();
  list.add(scRequest);
  StartContainersRequest allRequests =
      StartContainersRequest.newInstance(list);
  proxy.startContainers(allRequests);

  List<ContainerId> containerIds = new ArrayList<ContainerId>();
  containerIds.add(containerId);
  GetContainerStatusesRequest gcsRequest =
      GetContainerStatusesRequest.newInstance(containerIds);
  GetContainerStatusesResponse response =
      proxy.getContainerStatuses(gcsRequest);
  List<ContainerStatus> statuses = response.getContainerStatuses();

  //test remote exception
  boolean exception = false;
  try {
    StopContainersRequest stopRequest =
        recordFactory.newRecordInstance(StopContainersRequest.class);
    stopRequest.setContainerIds(containerIds);
    proxy.stopContainers(stopRequest);
    } catch (YarnException e) {
    exception = true;
    Assert.assertTrue(e.getMessage().contains(EXCEPTION_MSG));
    Assert.assertTrue(e.getMessage().contains(EXCEPTION_CAUSE));
    System.out.println("Test Exception is " + e.getMessage());
  } catch (Exception ex) {
    ex.printStackTrace();
  }
  Assert.assertTrue(exception);
  
  server.stop();
  Assert.assertNotNull(statuses.get(0));
  Assert.assertEquals(ContainerState.RUNNING, statuses.get(0).getState());
}