org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse Java Examples

@SuppressWarnings("deprecation")
@Test
public void testAllocateResponseWithoutIncDecContainers() {
  AllocateResponse r =
      AllocateResponse.newInstance(3, new ArrayList<ContainerStatus>(),
          new ArrayList<Container>(), new ArrayList<NodeReport>(), null,
          AMCommand.AM_RESYNC, 3, null, new ArrayList<NMToken>(), null, null);

  // serde
  AllocateResponseProto p = ((AllocateResponsePBImpl) r).getProto();
  r = new AllocateResponsePBImpl(p);

  // check value
  Assert.assertEquals(0, r.getIncreasedContainers().size());
  Assert.assertEquals(0, r.getDecreasedContainers().size());
}
 

public static AllocateResponse newAllocateResponse(int responseId,
    List<ContainerStatus> completedContainers,
    List<Container> allocatedContainers, List<NodeReport> updatedNodes,
    Resource availResources, AMCommand command, int numClusterNodes,
    PreemptionMessage preempt) {
  AllocateResponse response = recordFactory
      .newRecordInstance(AllocateResponse.class);
  response.setNumClusterNodes(numClusterNodes);
  response.setResponseId(responseId);
  response.setCompletedContainersStatuses(completedContainers);
  response.setAllocatedContainers(allocatedContainers);
  response.setUpdatedNodes(updatedNodes);
  response.setAvailableResources(availResources);
  response.setAMCommand(command);
  response.setPreemptionMessage(preempt);

  return response;
}
 

@Override
public AllocateResponse allocate(AllocateRequest request)
    throws YarnException, IOException {
  Assert.assertEquals("response ID mismatch",
      responseId, request.getResponseId());
  ++responseId;
  org.apache.hadoop.yarn.api.records.Token yarnToken = null;
  if (amToken != null) {
    yarnToken = org.apache.hadoop.yarn.api.records.Token.newInstance(
        amToken.getIdentifier(), amToken.getKind().toString(),
        amToken.getPassword(), amToken.getService().toString());
  }
  return AllocateResponse.newInstance(responseId,
      Collections.<ContainerStatus>emptyList(),
      Collections.<Container>emptyList(),
      Collections.<NodeReport>emptyList(),
      Resources.none(), null, 1, null,
      Collections.<NMToken>emptyList(),
      yarnToken,
      Collections.<ContainerResourceIncrease>emptyList(),
      Collections.<ContainerResourceDecrease>emptyList());
}
 

private int getAllocatedContainersNumber(
    AMRMClientImpl<ContainerRequest> amClient, int iterationsLeft)
    throws YarnException, IOException {
  int allocatedContainerCount = 0;
  while (iterationsLeft-- > 0) {
    Log.info(" == alloc " + allocatedContainerCount + " it left " + iterationsLeft);
    AllocateResponse allocResponse = amClient.allocate(0.1f);
    assertEquals(0, amClient.ask.size());
    assertEquals(0, amClient.release.size());
      
    assertEquals(nodeCount, amClient.getClusterNodeCount());
    allocatedContainerCount += allocResponse.getAllocatedContainers().size();
      
    if(allocatedContainerCount == 0) {
      // sleep to let NM's heartbeat to RM and trigger allocations
      sleep(100);
    }
  }
  return allocatedContainerCount;
}
 

public static AllocateResponse newAllocateResponse(int responseId,
    List<ContainerStatus> completedContainers,
    List<Container> allocatedContainers, List<NodeReport> updatedNodes,
    Resource availResources, AMCommand command, int numClusterNodes,
    PreemptionMessage preempt) {
  AllocateResponse response = recordFactory
      .newRecordInstance(AllocateResponse.class);
  response.setNumClusterNodes(numClusterNodes);
  response.setResponseId(responseId);
  response.setCompletedContainersStatuses(completedContainers);
  response.setAllocatedContainers(allocatedContainers);
  response.setUpdatedNodes(updatedNodes);
  response.setAvailableResources(availResources);
  response.setAMCommand(command);
  response.setPreemptionMessage(preempt);

  return response;
}
 

@Test(timeout = 3000000)
public void testRMIdentifierOnContainerAllocation() throws Exception {
  MockRM rm = new MockRM(conf);
  rm.start();

  // Register node1
  MockNM nm1 = rm.registerNode("127.0.0.1:1234", 6 * GB);

  // Submit an application
  RMApp app1 = rm.submitApp(2048);

  // kick the scheduling
  nm1.nodeHeartbeat(true);
  RMAppAttempt attempt1 = app1.getCurrentAppAttempt();
  MockAM am1 = rm.sendAMLaunched(attempt1.getAppAttemptId());
  am1.registerAppAttempt();

  am1.addRequests(new String[] { "127.0.0.1" }, GB, 1, 1);
  AllocateResponse alloc1Response = am1.schedule(); // send the request

  // kick the scheduler
  nm1.nodeHeartbeat(true);
  while (alloc1Response.getAllocatedContainers().size() < 1) {
    LOG.info("Waiting for containers to be created for app 1...");
    sleep(1000);
    alloc1Response = am1.schedule();
  }

  // assert RMIdentifer is set properly in allocated containers
  Container allocatedContainer =
      alloc1Response.getAllocatedContainers().get(0);
  ContainerTokenIdentifier tokenId =
      BuilderUtils.newContainerTokenIdentifier(allocatedContainer
        .getContainerToken());
  Assert.assertEquals(MockRM.getClusterTimeStamp(), tokenId.getRMIdentifier());
  rm.stop();
}
 

@Test (timeout = 5000)
public void testCallAMRMClientAsyncStopFromCallbackHandlerWithWaitFor()
    throws YarnException, IOException, InterruptedException {
  Configuration conf = new Configuration();
  final TestCallbackHandler2 callbackHandler = new TestCallbackHandler2();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);

  List<ContainerStatus> completed = Arrays.asList(
      ContainerStatus.newInstance(newContainerId(0, 0, 0, 0),
          ContainerState.COMPLETE, "", 0));
  final AllocateResponse response = createAllocateResponse(completed,
      new ArrayList<Container>(), null);

  when(client.allocate(anyFloat())).thenReturn(response);

  AMRMClientAsync<ContainerRequest> asyncClient =
      AMRMClientAsync.createAMRMClientAsync(client, 20, callbackHandler);
  callbackHandler.asynClient = asyncClient;
  asyncClient.init(conf);
  asyncClient.start();

  Supplier<Boolean> checker = new Supplier<Boolean>() {
    @Override
    public Boolean get() {
      return callbackHandler.notify;
    }
  };

  asyncClient.registerApplicationMaster("localhost", 1234, null);
  asyncClient.waitFor(checker);
  Assert.assertTrue(checker.get());
}
 

private AllocateResponse createAllocateResponse(
    List<ContainerStatus> completed, List<Container> allocated,
    List<NMToken> nmTokens) {
  AllocateResponse response =
      AllocateResponse.newInstance(0, completed, allocated,
          new ArrayList<NodeReport>(), null, null, 1, null, nmTokens);
  return response;
}
 

@SuppressWarnings("unchecked")
private void handleUpdatedNodes(AllocateResponse response) {
  // send event to the job about on updated nodes
  List<NodeReport> updatedNodes = response.getUpdatedNodes();
  if (!updatedNodes.isEmpty()) {

    // send event to the job to act upon completed tasks
    eventHandler.handle(new JobUpdatedNodesEvent(getJob().getID(),
        updatedNodes));

    // act upon running tasks
    HashSet<NodeId> unusableNodes = new HashSet<NodeId>();
    for (NodeReport nr : updatedNodes) {
      NodeState nodeState = nr.getNodeState();
      if (nodeState.isUnusable()) {
        unusableNodes.add(nr.getNodeId());
      }
    }
    for (int i = 0; i < 2; ++i) {
      HashMap<TaskAttemptId, Container> taskSet = i == 0 ? assignedRequests.maps
          : assignedRequests.reduces;
      // kill running containers
      for (Map.Entry<TaskAttemptId, Container> entry : taskSet.entrySet()) {
        TaskAttemptId tid = entry.getKey();
        NodeId taskAttemptNodeId = entry.getValue().getNodeId();
        if (unusableNodes.contains(taskAttemptNodeId)) {
          LOG.info("Killing taskAttempt:" + tid
              + " because it is running on unusable node:"
              + taskAttemptNodeId);
          eventHandler.handle(new TaskAttemptKillEvent(tid,
              "TaskAttempt killed because it ran on unusable node"
                  + taskAttemptNodeId));
        }
      }
    }
  }
}
 

public AllocateResponse allocate(
    List<ResourceRequest> resourceRequest, List<ContainerId> releases)
    throws Exception {
  final AllocateRequest req =
      AllocateRequest.newInstance(0, 0F, resourceRequest,
        releases, null);
  return allocate(req);
}
 

public AllocateResponse allocate(
    String host, int memory, int numContainers,
    List<ContainerId> releases, String labelExpression) throws Exception {
  List<ResourceRequest> reqs =
      createReq(new String[] { host }, memory, 1, numContainers,
          labelExpression);
  return allocate(reqs, releases);
}
 

/**
 * send out request for AM container
 */
protected void requestAMContainer()
        throws YarnException, IOException, InterruptedException {
  List<ResourceRequest> ask = new ArrayList<ResourceRequest>();
  ResourceRequest amRequest = createResourceRequest(
          BuilderUtils.newResource(MR_AM_CONTAINER_RESOURCE_MEMORY_MB,
                  MR_AM_CONTAINER_RESOURCE_VCORES),
          ResourceRequest.ANY, 1, 1);
  ask.add(amRequest);
  LOG.debug(MessageFormat.format("Application {0} sends out allocate " +
          "request for its AM", appId));
  final AllocateRequest request = this.createAllocateRequest(ask);

  UserGroupInformation ugi =
          UserGroupInformation.createRemoteUser(appAttemptId.toString());
  Token<AMRMTokenIdentifier> token = rm.getRMContext().getRMApps()
          .get(appAttemptId.getApplicationId())
          .getRMAppAttempt(appAttemptId).getAMRMToken();
  ugi.addTokenIdentifier(token.decodeIdentifier());
  AllocateResponse response = ugi.doAs(
          new PrivilegedExceptionAction<AllocateResponse>() {
    @Override
    public AllocateResponse run() throws Exception {
      return rm.getApplicationMasterService().allocate(request);
    }
  });
  if (response != null) {
    responseQueue.put(response);
  }
}
 

@Private
@VisibleForTesting
public AMRMClientAsyncImpl(AMRMClient<T> client, int intervalMs,
    CallbackHandler callbackHandler) {
  super(client, intervalMs, callbackHandler);
  heartbeatThread = new HeartbeatThread();
  handlerThread = new CallbackHandlerThread();
  responseQueue = new LinkedBlockingQueue<AllocateResponse>();
  keepRunning = true;
  savedException = null;
}
 

@SuppressWarnings("unchecked")
private void handleUpdatedNodes(AllocateResponse response) {
  // send event to the job about on updated nodes
  List<NodeReport> updatedNodes = response.getUpdatedNodes();
  if (!updatedNodes.isEmpty()) {

    // send event to the job to act upon completed tasks
    eventHandler.handle(new JobUpdatedNodesEvent(getJob().getID(),
        updatedNodes));

    // act upon running tasks
    HashSet<NodeId> unusableNodes = new HashSet<NodeId>();
    for (NodeReport nr : updatedNodes) {
      NodeState nodeState = nr.getNodeState();
      if (nodeState.isUnusable()) {
        unusableNodes.add(nr.getNodeId());
      }
    }
    for (int i = 0; i < 2; ++i) {
      HashMap<TaskAttemptId, Container> taskSet = i == 0 ? assignedRequests.maps
          : assignedRequests.reduces;
      // kill running containers
      for (Map.Entry<TaskAttemptId, Container> entry : taskSet.entrySet()) {
        TaskAttemptId tid = entry.getKey();
        NodeId taskAttemptNodeId = entry.getValue().getNodeId();
        if (unusableNodes.contains(taskAttemptNodeId)) {
          LOG.info("Killing taskAttempt:" + tid
              + " because it is running on unusable node:"
              + taskAttemptNodeId);
          eventHandler.handle(new TaskAttemptKillEvent(tid,
              "TaskAttempt killed because it ran on unusable node"
                  + taskAttemptNodeId));
        }
      }
    }
  }
}
 

private AllocateResponse allocate(ApplicationAttemptId attemptId,
    final AllocateRequest req) throws Exception {
  UserGroupInformation ugi =
      UserGroupInformation.createRemoteUser(attemptId.toString());
  org.apache.hadoop.security.token.Token<AMRMTokenIdentifier> token =
      rm.getRMContext().getRMApps().get(attemptId.getApplicationId())
        .getRMAppAttempt(attemptId).getAMRMToken();
  ugi.addTokenIdentifier(token.decodeIdentifier());
  return ugi.doAs(new PrivilegedExceptionAction<AllocateResponse>() {
    @Override
    public AllocateResponse run() throws Exception {
      return amService.allocate(req);
    }
  });
}
 

private AllocateResponse allocate(ApplicationAttemptId attemptId,
    final AllocateRequest req) throws Exception {
  UserGroupInformation ugi =
      UserGroupInformation.createRemoteUser(attemptId.toString());
  org.apache.hadoop.security.token.Token<AMRMTokenIdentifier> token =
      rm.getRMContext().getRMApps().get(attemptId.getApplicationId())
        .getRMAppAttempt(attemptId).getAMRMToken();
  ugi.addTokenIdentifier(token.decodeIdentifier());
  return ugi.doAs(new PrivilegedExceptionAction<AllocateResponse>() {
    @Override
    public AllocateResponse run() throws Exception {
      return amService.allocate(req);
    }
  });
}
 

@Test(timeout = 3000000)
public void testRMIdentifierOnContainerAllocation() throws Exception {
  MockRM rm = new MockRM(conf);
  rm.start();

  // Register node1
  MockNM nm1 = rm.registerNode("127.0.0.1:1234", 6 * GB);

  // Submit an application
  RMApp app1 = rm.submitApp(2048);

  // kick the scheduling
  nm1.nodeHeartbeat(true);
  RMAppAttempt attempt1 = app1.getCurrentAppAttempt();
  MockAM am1 = rm.sendAMLaunched(attempt1.getAppAttemptId());
  am1.registerAppAttempt();

  am1.addRequests(new String[] { "127.0.0.1" }, GB, 1, 1);
  AllocateResponse alloc1Response = am1.schedule(); // send the request

  // kick the scheduler
  nm1.nodeHeartbeat(true);
  while (alloc1Response.getAllocatedContainers().size() < 1) {
    LOG.info("Waiting for containers to be created for app 1...");
    sleep(1000);
    alloc1Response = am1.schedule();
  }

  // assert RMIdentifer is set properly in allocated containers
  Container allocatedContainer =
      alloc1Response.getAllocatedContainers().get(0);
  ContainerTokenIdentifier tokenId =
      BuilderUtils.newContainerTokenIdentifier(allocatedContainer
        .getContainerToken());
  Assert.assertEquals(MockRM.getClusterTimeStamp(), tokenId.getRMIdentifier());
  rm.stop();
}
 

@Test (timeout = 5000)
public void testCallAMRMClientAsyncStopFromCallbackHandler()
    throws YarnException, IOException, InterruptedException {
  Configuration conf = new Configuration();
  TestCallbackHandler2 callbackHandler = new TestCallbackHandler2();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);

  List<ContainerStatus> completed = Arrays.asList(
      ContainerStatus.newInstance(newContainerId(0, 0, 0, 0),
          ContainerState.COMPLETE, "", 0));
  final AllocateResponse response = createAllocateResponse(completed,
      new ArrayList<Container>(), null);

  when(client.allocate(anyFloat())).thenReturn(response);

  AMRMClientAsync<ContainerRequest> asyncClient =
      AMRMClientAsync.createAMRMClientAsync(client, 20, callbackHandler);
  callbackHandler.asynClient = asyncClient;
  asyncClient.init(conf);
  asyncClient.start();

  synchronized (callbackHandler.notifier) {
    asyncClient.registerApplicationMaster("localhost", 1234, null);
    while(callbackHandler.notify == false) {
      try {
        callbackHandler.notifier.wait();
      } catch (InterruptedException e) {
        e.printStackTrace();
      }
    }
  }
}
 

public AllocateResponse allocate(
    List<ResourceRequest> resourceRequest, List<ContainerId> releases)
    throws Exception {
  final AllocateRequest req =
      AllocateRequest.newInstance(0, 0F, resourceRequest,
        releases, null);
  return allocate(req);
}
 

@Override
public AllocateResponse allocate(AllocateRequest request)
    throws YarnException, IOException {
  AllocateRequestProto requestProto =
      ((AllocateRequestPBImpl) request).getProto();
  try {
    return new AllocateResponsePBImpl(proxy.allocate(null, requestProto));
  } catch (ServiceException e) {
    RPCUtil.unwrapAndThrowException(e);
    return null;
  }
}
 

@Test (timeout = 5000)
public void testCallAMRMClientAsyncStopFromCallbackHandlerWithWaitFor()
    throws YarnException, IOException, InterruptedException {
  Configuration conf = new Configuration();
  final TestCallbackHandler2 callbackHandler = new TestCallbackHandler2();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);

  List<ContainerStatus> completed = Arrays.asList(
      ContainerStatus.newInstance(newContainerId(0, 0, 0, 0),
          ContainerState.COMPLETE, "", 0));
  final AllocateResponse response = createAllocateResponse(completed,
      new ArrayList<Container>(), null);

  when(client.allocate(anyFloat())).thenReturn(response);

  AMRMClientAsync<ContainerRequest> asyncClient =
      AMRMClientAsync.createAMRMClientAsync(client, 20, callbackHandler);
  callbackHandler.asynClient = asyncClient;
  asyncClient.init(conf);
  asyncClient.start();

  Supplier<Boolean> checker = new Supplier<Boolean>() {
    @Override
    public Boolean get() {
      return callbackHandler.notify;
    }
  };

  asyncClient.registerApplicationMaster("localhost", 1234, null);
  asyncClient.waitFor(checker);
  Assert.assertTrue(checker.get());
}
 

protected void allocateContainersAndValidateNMTokens(MockAM am,
    ArrayList<Container> containersReceived, int totalContainerRequested,
    HashMap<String, Token> nmTokens, MockNM nm) throws Exception,
    InterruptedException {
  ArrayList<ContainerId> releaseContainerList = new ArrayList<ContainerId>();
  AllocateResponse response;
  ArrayList<ResourceRequest> resourceRequest =
      new ArrayList<ResourceRequest>();      
  while (containersReceived.size() < totalContainerRequested) {
    nm.nodeHeartbeat(true);
    LOG.info("requesting containers..");
    response =
        am.allocate(resourceRequest, releaseContainerList);
    containersReceived.addAll(response.getAllocatedContainers());
    if (!response.getNMTokens().isEmpty()) {
      for (NMToken nmToken : response.getNMTokens()) {
        String nodeId = nmToken.getNodeId().toString();
        if (nmTokens.containsKey(nodeId)) {
          Assert.fail("Duplicate NMToken received for : " + nodeId);
        }
        nmTokens.put(nodeId, nmToken.getToken());
      }
    }
    LOG.info("Got " + containersReceived.size()
        + " containers. Waiting to get " + totalContainerRequested);
    Thread.sleep(WAIT_SLEEP_MS);
  }
}
 

@Override
public AllocateResponse allocate(AllocateRequest request)
    throws YarnException, IOException {
  AllocateRequestProto requestProto =
      ((AllocateRequestPBImpl) request).getProto();
  try {
    return new AllocateResponsePBImpl(proxy.allocate(null, requestProto));
  } catch (ServiceException e) {
    RPCUtil.unwrapAndThrowException(e);
    return null;
  }
}
 

@Test(timeout = 15000)
public void testAllocateOnHA() throws YarnException, IOException {
  AllocateRequest request = AllocateRequest.newInstance(0, 50f,
      new ArrayList<ResourceRequest>(),
      new ArrayList<ContainerId>(),
      ResourceBlacklistRequest.newInstance(new ArrayList<String>(),
          new ArrayList<String>()));
  AllocateResponse response = amClient.allocate(request);
  Assert.assertEquals(response, this.cluster.createFakeAllocateResponse());
}
 

@Override
public AllocateResponse allocate(AllocateRequest request)
    throws YarnException, IOException {
  resetStartFailoverFlag(true);
  // make sure failover has been triggered
  Assert.assertTrue(waittingForFailOver());
  return createFakeAllocateResponse();
}
 

public void registerAppAttempt(ApplicationAttemptId attemptId) {
  AllocateResponse response =
      recordFactory.newRecordInstance(AllocateResponse.class);
  // set response id to -1 before application master for the following
  // attemptID get registered
  response.setResponseId(-1);
  LOG.info("Registering app attempt : " + attemptId);
  responseMap.put(attemptId, new AllocateResponseLock(response));
  rmContext.getNMTokenSecretManager().registerApplicationAttempt(attemptId);
}
 

public AllocateResponse doAllocateAs(UserGroupInformation ugi,
    final AllocateRequest req) throws Exception {
  req.setResponseId(++responseId);
  try {
    return ugi.doAs(new PrivilegedExceptionAction<AllocateResponse>() {
      @Override
      public AllocateResponse run() throws Exception {
        return amRMProtocol.allocate(req);
      }
    });
  } catch (UndeclaredThrowableException e) {
    throw (Exception) e.getCause();
  }
}
 

public AllocateResponse createFakeAllocateResponse() {
  return AllocateResponse.newInstance(-1,
      new ArrayList<ContainerStatus>(),
      new ArrayList<Container>(), new ArrayList<NodeReport>(),
      Resource.newInstance(1024, 2, 2), null, 1,
      null, new ArrayList<NMToken>());
}
 

@Private
@VisibleForTesting
public AMRMClientAsyncImpl(AMRMClient<T> client, int intervalMs,
    CallbackHandler callbackHandler) {
  super(client, intervalMs, callbackHandler);
  heartbeatThread = new HeartbeatThread();
  handlerThread = new CallbackHandlerThread();
  responseQueue = new LinkedBlockingQueue<AllocateResponse>();
  keepRunning = true;
  savedException = null;
}
 

public void registerAppAttempt(ApplicationAttemptId attemptId) {
  AllocateResponse response =
      recordFactory.newRecordInstance(AllocateResponse.class);
  // set response id to -1 before application master for the following
  // attemptID get registered
  response.setResponseId(-1);
  LOG.info("Registering app attempt : " + attemptId);
  responseMap.put(attemptId, new AllocateResponseLock(response));
  rmContext.getNMTokenSecretManager().registerApplicationAttempt(attemptId);
}