Java Code Examples for org.apache.hadoop.yarn.api.records.Resource#newInstance()

/**
 * This method produces an Allocation that includes the current view
 * of the resources that will be allocated to and preempted from this
 * application.
 *
 * @param rc
 * @param clusterResource
 * @param minimumAllocation
 * @return an allocation
 */
public synchronized Allocation getAllocation(ResourceCalculator rc,
    Resource clusterResource, Resource minimumAllocation) {

  Set<ContainerId> currentContPreemption = Collections.unmodifiableSet(
      new HashSet<ContainerId>(containersToPreempt));
  containersToPreempt.clear();
  Resource tot = Resource.newInstance(0, 0, 0);
  for(ContainerId c : currentContPreemption){
    Resources.addTo(tot,
        liveContainers.get(c).getContainer().getResource());
  }
  int numCont = (int) Math.ceil(
      Resources.divide(rc, clusterResource, tot, minimumAllocation));
  ResourceRequest rr = ResourceRequest.newInstance(
      Priority.UNDEFINED, ResourceRequest.ANY,
      minimumAllocation, numCont);
  ContainersAndNMTokensAllocation allocation =
      pullNewlyAllocatedContainersAndNMTokens();
  Resource headroom = getHeadroom();
  setApplicationHeadroomForMetrics(headroom);
  return new Allocation(allocation.getContainerList(), headroom, null,
    currentContPreemption, Collections.singletonList(rr),
    allocation.getNMTokenList());
}
 

TempQueue(String queueName, Resource current, Resource pending,
    Resource guaranteed, Resource maxCapacity, float guaranteedRatio,boolean preemptionDisabled) {
  this.queueName = queueName;
  this.current = current;
  this.pending = pending;
  this.guaranteed = guaranteed;
  this.maxCapacity = maxCapacity;
  this.idealAssigned = Resource.newInstance(0, 0);
  this.actuallyPreempted = Resource.newInstance(0, 0);
  this.toBePreempted = Resource.newInstance(0, 0);
  this.normalizedGuarantee = Float.NaN;
  this.children = new ArrayList<TempQueue>();
  //both leaf and parent node may have untouchable resource
  this.untouchableExtra = Resource.newInstance(0, 0);
  //only leaf node has preemptable extra
  this.preemptableExtra = Resource.newInstance(0, 0);
  this.guaranteedRatio  = guaranteedRatio;
  this.preemptionDisabled = preemptionDisabled;
}
 

@Test
public void testResourceIncreaseContext() {
  byte[] identifier = new byte[] { 1, 2, 3, 4 };
  Token token = Token.newInstance(identifier, "", "".getBytes(), "");
  ContainerId containerId = ContainerId
      .newContainerId(ApplicationAttemptId.newInstance(
          ApplicationId.newInstance(1234, 3), 3), 7);
  Resource resource = Resource.newInstance(1023, 3, 4);
  ContainerResourceIncrease ctx = ContainerResourceIncrease.newInstance(
      containerId, resource, token);

  // get proto and recover to ctx
  ContainerResourceIncreaseProto proto = 
      ((ContainerResourceIncreasePBImpl) ctx).getProto();
  ctx = new ContainerResourceIncreasePBImpl(proto);

  // check values
  Assert.assertEquals(ctx.getCapability(), resource);
  Assert.assertEquals(ctx.getContainerId(), containerId);
  Assert.assertTrue(Arrays.equals(ctx.getContainerToken().getIdentifier()
      .array(), identifier));
}
 

@Test
public void testDefaultTwoResourceTypeWithYarnSupport() {
	assumeTrue(HadoopUtils.isMinHadoopVersion(2, 10));

	final Resource resource = Resource.newInstance(100, 1);

	// make sure that Resource has at least two associated resources (cpu and memory)
	final Map<String, Long> resourcesResult = ResourceInformationReflector.INSTANCE.getAllResourceInfos(resource);
	assertThat(resourcesResult.size(), is(2));
}
 

@Override
public void validate(Plan plan, ReservationAllocation reservation)
    throws PlanningException {

  ReservationAllocation oldReservation =
      plan.getReservationById(reservation.getReservationId());

  // check updates are using same name
  if (oldReservation != null
      && !oldReservation.getUser().equals(reservation.getUser())) {
    throw new MismatchedUserException(
        "Updating an existing reservation with mismatching user:"
            + oldReservation.getUser() + " != " + reservation.getUser());
  }

  long startTime = reservation.getStartTime();
  long endTime = reservation.getEndTime();
  long step = plan.getStep();

  // for every instant in time, check we are respecting cluster capacity
  for (long t = startTime; t < endTime; t += step) {
    Resource currExistingAllocTot = plan.getTotalCommittedResources(t);
    Resource currNewAlloc = reservation.getResourcesAtTime(t);
    Resource currOldAlloc = Resource.newInstance(0, 0);
    if (oldReservation != null) {
      oldReservation.getResourcesAtTime(t);
    }
    // check the cluster is never over committed
    // currExistingAllocTot + currNewAlloc - currOldAlloc >
    // capPlan.getTotalCapacity()
    if (Resources.greaterThan(plan.getResourceCalculator(), plan
        .getTotalCapacity(), Resources.subtract(
        Resources.add(currExistingAllocTot, currNewAlloc), currOldAlloc),
        plan.getTotalCapacity())) {
      throw new ResourceOverCommitException("Resources at time " + t
          + " would be overcommitted by " + "accepting reservation: "
          + reservation.getReservationId());
    }
  }
}
 

@Override
public void validate(Plan plan, ReservationAllocation reservation)
    throws PlanningException {

  ReservationAllocation oldReservation =
      plan.getReservationById(reservation.getReservationId());

  // check updates are using same name
  if (oldReservation != null
      && !oldReservation.getUser().equals(reservation.getUser())) {
    throw new MismatchedUserException(
        "Updating an existing reservation with mismatching user:"
            + oldReservation.getUser() + " != " + reservation.getUser());
  }

  long startTime = reservation.getStartTime();
  long endTime = reservation.getEndTime();
  long step = plan.getStep();

  // for every instant in time, check we are respecting cluster capacity
  for (long t = startTime; t < endTime; t += step) {
    Resource currExistingAllocTot = plan.getTotalCommittedResources(t);
    Resource currNewAlloc = reservation.getResourcesAtTime(t);
    Resource currOldAlloc = Resource.newInstance(0, 0, 0);
    if (oldReservation != null) {
      oldReservation.getResourcesAtTime(t);
    }
    // check the cluster is never over committed
    // currExistingAllocTot + currNewAlloc - currOldAlloc >
    // capPlan.getTotalCapacity()
    if (Resources.greaterThan(plan.getResourceCalculator(), plan
        .getTotalCapacity(), Resources.subtract(
        Resources.add(currExistingAllocTot, currNewAlloc), currOldAlloc),
        plan.getTotalCapacity())) {
      throw new ResourceOverCommitException("Resources at time " + t
          + " would be overcommitted by " + "accepting reservation: "
          + reservation.getReservationId());
    }
  }
}
 

@Override
protected void serviceInit(Configuration conf) throws Exception {
  int memoryMb = 
      conf.getInt(
          YarnConfiguration.NM_PMEM_MB, YarnConfiguration.DEFAULT_NM_PMEM_MB);
  float vMemToPMem =             
      conf.getFloat(
          YarnConfiguration.NM_VMEM_PMEM_RATIO, 
          YarnConfiguration.DEFAULT_NM_VMEM_PMEM_RATIO); 
  int virtualMemoryMb = (int)Math.ceil(memoryMb * vMemToPMem);
  
  int virtualCores =
      conf.getInt(
          YarnConfiguration.NM_VCORES, YarnConfiguration.DEFAULT_NM_VCORES);

  int gpuCores =
      conf.getInt(
          YarnConfiguration.NM_GCORES, YarnConfiguration.DEFAULT_NM_GCORES);

  this.totalResource = Resource.newInstance(memoryMb, virtualCores, gpuCores);
  metrics.addResource(totalResource);
  this.tokenKeepAliveEnabled = isTokenKeepAliveEnabled(conf);
  this.tokenRemovalDelayMs =
      conf.getInt(YarnConfiguration.RM_NM_EXPIRY_INTERVAL_MS,
          YarnConfiguration.DEFAULT_RM_NM_EXPIRY_INTERVAL_MS);

  this.minimumResourceManagerVersion = conf.get(
      YarnConfiguration.NM_RESOURCEMANAGER_MINIMUM_VERSION,
      YarnConfiguration.DEFAULT_NM_RESOURCEMANAGER_MINIMUM_VERSION);
  
  // Default duration to track stopped containers on nodemanager is 10Min.
  // This should not be assigned very large value as it will remember all the
  // containers stopped during that time.
  durationToTrackStoppedContainers =
      conf.getLong(YARN_NODEMANAGER_DURATION_TO_TRACK_STOPPED_CONTAINERS,
        600000);
  if (durationToTrackStoppedContainers < 0) {
    String message = "Invalid configuration for "
      + YARN_NODEMANAGER_DURATION_TO_TRACK_STOPPED_CONTAINERS + " default "
        + "value is 10Min(600000).";
    LOG.error(message);
    throw new YarnException(message);
  }
  if (LOG.isDebugEnabled()) {
    LOG.debug(YARN_NODEMANAGER_DURATION_TO_TRACK_STOPPED_CONTAINERS + " :"
      + durationToTrackStoppedContainers);
  }
  super.serviceInit(conf);
  LOG.info("Initialized nodemanager for " + nodeId + ":" +
      " physical-memory=" + memoryMb + " virtual-memory=" + virtualMemoryMb +
      " virtual-cores=" + virtualCores + " gpu-cores=" + gpuCores);
}
 

/**
 * This tests whether a containerId is serialized/deserialized with epoch.
 *
 * @throws IOException
 * @throws InterruptedException
 * @throws YarnException
 */
private void testContainerTokenWithEpoch(Configuration conf)
    throws IOException, InterruptedException, YarnException {

  LOG.info("Running test for serializing/deserializing containerIds");

  NMTokenSecretManagerInRM nmTokenSecretManagerInRM =
      yarnCluster.getResourceManager().getRMContext()
          .getNMTokenSecretManager();
  ApplicationId appId = ApplicationId.newInstance(1, 1);
  ApplicationAttemptId appAttemptId =
      ApplicationAttemptId.newInstance(appId, 0);
  ContainerId cId = ContainerId.newContainerId(appAttemptId, (5L << 40) | 3L);
  NodeManager nm = yarnCluster.getNodeManager(0);
  NMTokenSecretManagerInNM nmTokenSecretManagerInNM =
      nm.getNMContext().getNMTokenSecretManager();
  String user = "test";

  waitForNMToReceiveNMTokenKey(nmTokenSecretManagerInNM, nm);

  NodeId nodeId = nm.getNMContext().getNodeId();

  // Both id should be equal.
  Assert.assertEquals(nmTokenSecretManagerInNM.getCurrentKey().getKeyId(),
      nmTokenSecretManagerInRM.getCurrentKey().getKeyId());

  // Creating a normal Container Token
  RMContainerTokenSecretManager containerTokenSecretManager =
      yarnCluster.getResourceManager().getRMContext().
          getContainerTokenSecretManager();
  Resource r = Resource.newInstance(1230, 2);
  Token containerToken =
      containerTokenSecretManager.createContainerToken(cId, nodeId, user, r,
          Priority.newInstance(0), 0);
  
  ContainerTokenIdentifier containerTokenIdentifier =
      new ContainerTokenIdentifier();
  byte[] tokenIdentifierContent = containerToken.getIdentifier().array();
  DataInputBuffer dib = new DataInputBuffer();
  dib.reset(tokenIdentifierContent, tokenIdentifierContent.length);
  containerTokenIdentifier.readFields(dib);
  
  
  Assert.assertEquals(cId, containerTokenIdentifier.getContainerID());
  Assert.assertEquals(
      cId.toString(), containerTokenIdentifier.getContainerID().toString());

  Token nmToken =
      nmTokenSecretManagerInRM.createNMToken(appAttemptId, nodeId, user);

  YarnRPC rpc = YarnRPC.create(conf);
  testStartContainer(rpc, appAttemptId, nodeId, containerToken, nmToken,
      false);

  List<ContainerId> containerIds = new LinkedList<ContainerId>();
  containerIds.add(cId);
  ContainerManagementProtocol proxy
      = getContainerManagementProtocolProxy(rpc, nmToken, nodeId, user);
  GetContainerStatusesResponse res = proxy.getContainerStatuses(
      GetContainerStatusesRequest.newInstance(containerIds));
  Assert.assertNotNull(res.getContainerStatuses().get(0));
  Assert.assertEquals(
      cId, res.getContainerStatuses().get(0).getContainerId());
  Assert.assertEquals(cId.toString(),
      res.getContainerStatuses().get(0).getContainerId().toString());
}
 

/**
 * This tests a malice user getting a proper token but then messing with it by
 * tampering with containerID/Resource etc.. His/her containers should be
 * rejected.
 * 
 * @throws IOException
 * @throws InterruptedException
 * @throws YarnException
 */
private void testContainerToken(Configuration conf) throws IOException,
    InterruptedException, YarnException {

  LOG.info("Running test for malice user");
  /*
   * We need to check for containerToken (authorization).
   * Here we will be assuming that we have valid NMToken  
   * 1) ContainerToken used is expired.
   * 2) ContainerToken is tampered (resource is modified).
   */
  NMTokenSecretManagerInRM nmTokenSecretManagerInRM =
      yarnCluster.getResourceManager().getRMContext()
        .getNMTokenSecretManager();
  ApplicationId appId = ApplicationId.newInstance(1, 1);
  ApplicationAttemptId appAttemptId =
      ApplicationAttemptId.newInstance(appId, 0);
  ContainerId cId = ContainerId.newContainerId(appAttemptId, 0);
  NodeManager nm = yarnCluster.getNodeManager(0);
  NMTokenSecretManagerInNM nmTokenSecretManagerInNM =
      nm.getNMContext().getNMTokenSecretManager();
  String user = "test";
  
  waitForNMToReceiveNMTokenKey(nmTokenSecretManagerInNM, nm);

  NodeId nodeId = nm.getNMContext().getNodeId();
  
  // Both id should be equal.
  Assert.assertEquals(nmTokenSecretManagerInNM.getCurrentKey().getKeyId(),
      nmTokenSecretManagerInRM.getCurrentKey().getKeyId());
  
  
  RMContainerTokenSecretManager containerTokenSecretManager =
      yarnCluster.getResourceManager().getRMContext().
          getContainerTokenSecretManager();
  
  Resource r = Resource.newInstance(1230, 2, 2);
  
  Token containerToken = 
      containerTokenSecretManager.createContainerToken(
          cId, nodeId, user, r, Priority.newInstance(0), 0);
  
  ContainerTokenIdentifier containerTokenIdentifier = 
      getContainerTokenIdentifierFromToken(containerToken);
  
  // Verify new compatible version ContainerTokenIdentifier can work successfully.
  ContainerTokenIdentifierForTest newVersionTokenIdentifier = 
      new ContainerTokenIdentifierForTest(containerTokenIdentifier, "message");
  byte[] password = 
      containerTokenSecretManager.createPassword(newVersionTokenIdentifier);
  
  Token newContainerToken = BuilderUtils.newContainerToken(
      nodeId, password, newVersionTokenIdentifier);
  
  Token nmToken =
          nmTokenSecretManagerInRM.createNMToken(appAttemptId, nodeId, user);
  YarnRPC rpc = YarnRPC.create(conf);
  Assert.assertTrue(testStartContainer(rpc, appAttemptId, nodeId,
      newContainerToken, nmToken, false).isEmpty());
  
  // Creating a tampered Container Token
  RMContainerTokenSecretManager tamperedContainerTokenSecretManager =
      new RMContainerTokenSecretManager(conf);
  tamperedContainerTokenSecretManager.rollMasterKey();
  do {
    tamperedContainerTokenSecretManager.rollMasterKey();
    tamperedContainerTokenSecretManager.activateNextMasterKey();
  } while (containerTokenSecretManager.getCurrentKey().getKeyId()
      == tamperedContainerTokenSecretManager.getCurrentKey().getKeyId());
  
  ContainerId cId2 = ContainerId.newContainerId(appAttemptId, 1);
  // Creating modified containerToken
  Token containerToken2 =
      tamperedContainerTokenSecretManager.createContainerToken(cId2, nodeId,
          user, r, Priority.newInstance(0), 0);
  
  StringBuilder sb = new StringBuilder("Given Container ");
  sb.append(cId2);
  sb.append(" seems to have an illegally generated token.");
  Assert.assertTrue(testStartContainer(rpc, appAttemptId, nodeId,
      containerToken2, nmToken, true).contains(sb.toString()));
}
 

@Override
public Resource getAvailableResources() {
  return Resource.newInstance(2048, 2);
}
 

/**
 * Checks if any containers were allocated in the current scheduler run and
 * launches the corresponding Mesos tasks. It also updates the node
 * capacity depending on what portion of the consumed offers were actually
 * used.
 */
@VisibleForTesting
protected void handleContainerAllocation(RMNode rmNode) {
  String host = rmNode.getNodeID().getHost();

  ConsumedOffer consumedOffer = offerLifecycleMgr.drainConsumedOffer(host);
  if (consumedOffer == null) {
    LOGGER.debug("No offer consumed for {}", host);
    return;
  }

  Node node = nodeStore.getNode(host);
  Set<RMContainer> containersBeforeSched = node.getContainerSnapshot();
  Set<RMContainer> containersAfterSched = new HashSet<>(node.getNode().getRunningContainers());

  Set<RMContainer> containersAllocatedByMesosOffer = (containersBeforeSched == null) ? containersAfterSched : Sets.difference(
      containersAfterSched, containersBeforeSched);

  if (containersAllocatedByMesosOffer.isEmpty()) {
    LOGGER.debug("No containers allocated using Mesos offers for host: {}", host);
    for (Protos.Offer offer : consumedOffer.getOffers()) {
      offerLifecycleMgr.declineOffer(offer);
    }
    decrementNodeCapacity(rmNode, OfferUtils.getYarnResourcesFromMesosOffers(consumedOffer.getOffers()));
  } else {
    LOGGER.debug("Containers allocated using Mesos offers for host: {} count: {}", host, containersAllocatedByMesosOffer.size());

    // Identify the Mesos tasks that need to be launched
    List<Protos.TaskInfo> tasks = Lists.newArrayList();
    Resource resUsed = Resource.newInstance(0, 0);

    for (RMContainer newContainer : containersAllocatedByMesosOffer) {
      tasks.add(getTaskInfoForContainer(newContainer, consumedOffer, node));
      resUsed = Resources.add(resUsed, newContainer.getAllocatedResource());
    }

    // Reduce node capacity to account for unused offers
    Resource resOffered = OfferUtils.getYarnResourcesFromMesosOffers(consumedOffer.getOffers());
    Resource resUnused = Resources.subtract(resOffered, resUsed);
    decrementNodeCapacity(rmNode, resUnused);
    myriadDriver.getDriver().launchTasks(consumedOffer.getOfferIds(), tasks);
  }

  // No need to hold on to the snapshot anymore
  node.removeContainerSnapshot();
}
 

/**
 * Removes a resource for the specified interval
 * 
 * @param reservationInterval the interval for which the resource is to be
 *          removed
 * @param capacity the resource to be removed
 * @return true if removal is successful, false otherwise
 */
public boolean removeInterval(ReservationInterval reservationInterval,
    ReservationRequest capacity) {
  Resource totCap =
      Resources.multiply(capacity.getCapability(),
          (float) capacity.getNumContainers());
  if (totCap.equals(ZERO_RESOURCE)) {
    return true;
  }
  writeLock.lock();
  try {
    long startKey = reservationInterval.getStartTime();
    long endKey = reservationInterval.getEndTime();
    // update the start key
    NavigableMap<Long, Resource> ticks =
        cumulativeCapacity.headMap(endKey, false);
    // Decrease all the capacities of overlapping intervals
    SortedMap<Long, Resource> overlapSet = ticks.tailMap(startKey);
    if (overlapSet != null && !overlapSet.isEmpty()) {
      Resource updatedCapacity = Resource.newInstance(0, 0);
      long currentKey = -1;
      for (Iterator<Entry<Long, Resource>> overlapEntries =
          overlapSet.entrySet().iterator(); overlapEntries.hasNext();) {
        Entry<Long, Resource> entry = overlapEntries.next();
        currentKey = entry.getKey();
        updatedCapacity = Resources.subtract(entry.getValue(), totCap);
        // update each entry between start and end key
        cumulativeCapacity.put(currentKey, updatedCapacity);
      }
      // Remove the first overlap entry if it is same as previous after
      // updation
      Long firstKey = overlapSet.firstKey();
      if (isSameAsPrevious(firstKey, overlapSet.get(firstKey))) {
        cumulativeCapacity.remove(firstKey);
      }
      // Remove the next entry if it is same as end entry after updation
      if ((currentKey != -1) && (isSameAsNext(currentKey, updatedCapacity))) {
        cumulativeCapacity.remove(cumulativeCapacity.higherKey(currentKey));
      }
    }
    return true;
  } finally {
    writeLock.unlock();
  }
}
 

public static Resource createResourceRequestFromTaskConfig(
    PlanTaskConfiguration taskConfig) {
  return Resource.newInstance(taskConfig.getMemoryMb(), taskConfig.getVirtualCores());
}
 

public static RMNode getRMNode(String host, int port, int memory, int cores) {
  Resource resource = Resource.newInstance(memory, cores);
  return getRMNode(host, port, resource);
}
 

@Test(timeout = 10000l)
public void testAssignmentOnShutdown()
    throws IOException, InterruptedException, ExecutionException {
  LOG.info("Test testAssignmentOnShutdown");
  Configuration tezConf = new Configuration(new YarnConfiguration());
  tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false);
  tezConf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
  tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 0);
  tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 0);
  tezConf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 0);

  CapturingEventHandler eventHandler = new CapturingEventHandler();
  TezDAGID dagID = TezDAGID.getInstance("0", 0, 0);

  AMRMClient<CookieContainerRequest> rmClientCore = new AMRMClientForTest();
  TezAMRMClientAsync<CookieContainerRequest> rmClient = spy(new AMRMClientAsyncForTest(rmClientCore, 100));

  AppContext appContext = mock(AppContext.class);
  doReturn(new Configuration(false)).when(appContext).getAMConf();
  AMContainerMap amContainerMap = new AMContainerMap(mock(ContainerHeartbeatHandler.class),
      mock(TaskCommunicatorManagerInterface.class), new ContainerContextMatcher(), appContext);
  AMNodeTracker amNodeTracker = new AMNodeTracker(eventHandler, appContext);
  doReturn(amContainerMap).when(appContext).getAllContainers();
  doReturn(amNodeTracker).when(appContext).getNodeTracker();
  doReturn(DAGAppMasterState.SUCCEEDED).when(appContext).getAMState();
  doReturn(true).when(appContext).isAMInCompletionState();
  doReturn(dagID).when(appContext).getCurrentDAGID();
  doReturn(mock(ClusterInfo.class)).when(appContext).getClusterInfo();

  TaskSchedulerManager taskSchedulerManagerReal =
      new TaskSchedulerManagerForTest(appContext, eventHandler, rmClient,
          new AlwaysMatchesContainerMatcher(), TezUtils.createUserPayloadFromConf(tezConf));
  TaskSchedulerManager taskSchedulerManager = spy(taskSchedulerManagerReal);
  taskSchedulerManager.init(tezConf);
  taskSchedulerManager.start();

  TaskSchedulerWithDrainableContext taskScheduler = (TaskSchedulerWithDrainableContext)
      ((TaskSchedulerManagerForTest) taskSchedulerManager)
      .getSpyTaskScheduler();
  TaskSchedulerContextDrainable drainableAppCallback = taskScheduler.getDrainableAppCallback();
  AtomicBoolean drainNotifier = new AtomicBoolean(false);
  taskScheduler.delayedContainerManager.drainedDelayedContainersForTest = drainNotifier;

  Resource resource1 = Resource.newInstance(1024, 1);
  String[] host1 = {"host1"};

  String []racks = {"/default-rack"};
  Priority priority1 = Priority.newInstance(1);

  TezVertexID vertexID1 = TezVertexID.getInstance(dagID, 1);

  //Vertex 1, Task 1, Attempt 1, host1
  TezTaskAttemptID taID11 = TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexID1, 1), 1);
  TaskAttempt ta11 = mock(TaskAttempt.class);
  AMSchedulerEventTALaunchRequest lrEvent1 = createLaunchRequestEvent(taID11, ta11, resource1,
      host1, racks, priority1);
  taskSchedulerManager.handleEvent(lrEvent1);

  Container container1 = createContainer(1, "host1", resource1, priority1);

  // One container allocated.
  taskScheduler.onContainersAllocated(Collections.singletonList(container1));
  drainableAppCallback.drain();
  verify(taskSchedulerManager, times(0)).taskAllocated(eq(0), eq(ta11),
      any(Object.class), eq(container1));
  taskScheduler.shutdown();
  taskSchedulerManager.close();
}
 

@Test (timeout=60000)
public void testAMRMClientMatchingFitInferredRack() throws YarnException, IOException {
  AMRMClientImpl<ContainerRequest> amClient = null;
  try {
    // start am rm client
    amClient = new AMRMClientImpl<ContainerRequest>();
    amClient.init(conf);
    amClient.start();
    amClient.registerApplicationMaster("Host", 10000, "");
    
    Resource capability = Resource.newInstance(1024, 2);

    ContainerRequest storedContainer1 = 
        new ContainerRequest(capability, nodes, null, priority);
    amClient.addContainerRequest(storedContainer1);

    // verify matching with original node and inferred rack
    List<? extends Collection<ContainerRequest>> matches;
    ContainerRequest storedRequest;
    // exact match node
    matches = amClient.getMatchingRequests(priority, node, capability);
    verifyMatches(matches, 1);
    storedRequest = matches.get(0).iterator().next();
    assertEquals(storedContainer1, storedRequest);
    // inferred match rack
    matches = amClient.getMatchingRequests(priority, rack, capability);
    verifyMatches(matches, 1);
    storedRequest = matches.get(0).iterator().next();
    assertEquals(storedContainer1, storedRequest);
    
    // inferred rack match no longer valid after request is removed
    amClient.removeContainerRequest(storedContainer1);
    matches = amClient.getMatchingRequests(priority, rack, capability);
    assertTrue(matches.isEmpty());
    
    amClient.unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED,
        null, null);

  } finally {
    if (amClient != null && amClient.getServiceState() == STATE.STARTED) {
      amClient.stop();
    }
  }
}
 

@Override
public RMAppMetrics getRMAppMetrics() {
  return new RMAppMetrics(Resource.newInstance(0, 0, 0), 0, 0, 0, 0, 0);
}
 

@Override
protected void serviceInit(Configuration conf) throws Exception {
  int memoryMb = 
      conf.getInt(
          YarnConfiguration.NM_PMEM_MB, YarnConfiguration.DEFAULT_NM_PMEM_MB);
  float vMemToPMem =             
      conf.getFloat(
          YarnConfiguration.NM_VMEM_PMEM_RATIO, 
          YarnConfiguration.DEFAULT_NM_VMEM_PMEM_RATIO); 
  int virtualMemoryMb = (int)Math.ceil(memoryMb * vMemToPMem);
  
  int virtualCores =
      conf.getInt(
          YarnConfiguration.NM_VCORES, YarnConfiguration.DEFAULT_NM_VCORES);

  this.totalResource = Resource.newInstance(memoryMb, virtualCores);
  metrics.addResource(totalResource);
  this.tokenKeepAliveEnabled = isTokenKeepAliveEnabled(conf);
  this.tokenRemovalDelayMs =
      conf.getInt(YarnConfiguration.RM_NM_EXPIRY_INTERVAL_MS,
          YarnConfiguration.DEFAULT_RM_NM_EXPIRY_INTERVAL_MS);

  this.minimumResourceManagerVersion = conf.get(
      YarnConfiguration.NM_RESOURCEMANAGER_MINIMUM_VERSION,
      YarnConfiguration.DEFAULT_NM_RESOURCEMANAGER_MINIMUM_VERSION);
  
  // Default duration to track stopped containers on nodemanager is 10Min.
  // This should not be assigned very large value as it will remember all the
  // containers stopped during that time.
  durationToTrackStoppedContainers =
      conf.getLong(YARN_NODEMANAGER_DURATION_TO_TRACK_STOPPED_CONTAINERS,
        600000);
  if (durationToTrackStoppedContainers < 0) {
    String message = "Invalid configuration for "
      + YARN_NODEMANAGER_DURATION_TO_TRACK_STOPPED_CONTAINERS + " default "
        + "value is 10Min(600000).";
    LOG.error(message);
    throw new YarnException(message);
  }
  if (LOG.isDebugEnabled()) {
    LOG.debug(YARN_NODEMANAGER_DURATION_TO_TRACK_STOPPED_CONTAINERS + " :"
      + durationToTrackStoppedContainers);
  }
  super.serviceInit(conf);
  LOG.info("Initialized nodemanager for " + nodeId + ":" +
      " physical-memory=" + memoryMb + " virtual-memory=" + virtualMemoryMb +
      " virtual-cores=" + virtualCores);
}
 

@Test
public void testContainerExpired() throws Exception {
  TezAMRMClientAsync<CookieContainerRequest> mockRMClient = spy(
      new AMRMClientAsyncForTest(new AMRMClientForTest(), 100));

  Configuration conf = new Configuration();
  // to match all in the same pass
  conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 0);
  // to release immediately after deallocate
  conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 0);
  conf.setLong(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 0);

  TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(DEFAULT_APP_HOST, DEFAULT_APP_PORT, DEFAULT_APP_URL, conf);
  final TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);
  TaskSchedulerWithDrainableContext scheduler =
      new TaskSchedulerWithDrainableContext(drainableAppCallback, mockRMClient);

  scheduler.initialize();
  scheduler.start();
  drainableAppCallback.drain();

  Object mockTask1 = new MockTask("task1");
  Object mockCookie1 = new Object();
  Resource mockCapability = Resource.newInstance(1024, 1);
  String[] hosts = {"host1", "host5"};
  String[] racks = {"/default-rack", "/default-rack"};
  final Priority mockPriority1 = Priority.newInstance(1);
  final Priority mockPriority2 = Priority.newInstance(2);
  Object mockTask2 = new MockTask("task2");
  Object mockCookie2 = new Object();
  ArgumentCaptor<CookieContainerRequest> requestCaptor =
      ArgumentCaptor.forClass(CookieContainerRequest.class);

  scheduler.allocateTask(mockTask2, mockCapability, hosts,
      racks, mockPriority2, null, mockCookie2);
  drainableAppCallback.drain();
  verify(mockRMClient, times(1)).
                              addContainerRequest(requestCaptor.capture());
  CookieContainerRequest request2 = requestCaptor.getValue();

  scheduler.allocateTask(mockTask1, mockCapability, hosts,
      racks, mockPriority1, null, mockCookie1);
  drainableAppCallback.drain();
  verify(mockRMClient, times(2)).
                              addContainerRequest(requestCaptor.capture());

  List<Container> containers = new ArrayList<Container>();
  // sending only lower priority container to make sure its not matched
  NodeId host2 = NodeId.newInstance("host2", 2);
  ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
  ContainerId mockCId2 = ContainerId.newContainerId(attemptId, 2);
  Container mockContainer2 = Container.newInstance(mockCId2, host2, null, mockCapability, mockPriority2, null);
  containers.add(mockContainer2);

  scheduler.onContainersAllocated(containers);

  List<ContainerStatus> statuses = new ArrayList<ContainerStatus>();
  ContainerStatus mockStatus2 = mock(ContainerStatus.class);
  when(mockStatus2.getContainerId()).thenReturn(mockCId2);
  statuses.add(mockStatus2);
  scheduler.onContainersCompleted(statuses);

  verify(mockApp, times(0)).taskAllocated(any(), any(), any(Container.class));
  verify(mockRMClient, times(3)).addContainerRequest(requestCaptor.capture());
  CookieContainerRequest resubmitRequest = requestCaptor.getValue();
  assertEquals(request2.getCookie().getTask(), resubmitRequest.getCookie().getTask());
  assertEquals(request2.getCookie().getAppCookie(), resubmitRequest.getCookie().getAppCookie());
  assertEquals(request2.getCookie().getContainerSignature(), resubmitRequest.getCookie().getContainerSignature());
  assertEquals(request2.getCapability(), resubmitRequest.getCapability());
  assertEquals(request2.getPriority(), resubmitRequest.getPriority());

  // verify container is not re-requested when nothing at that priority
  assertFalse(scheduler.deallocateTask(mockTask2, true, null, null));
  scheduler.onContainersAllocated(containers);
  scheduler.onContainersCompleted(statuses);
  verify(mockApp, times(0)).taskAllocated(any(), any(), any(Container.class));
  verify(mockRMClient, times(3)).addContainerRequest(requestCaptor.capture());
}
 

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());
}