Java Code Examples for org.apache.hadoop.yarn.client.api.AMRMClient#ContainerRequest

/**
 * @throws Exception If failed.
 */
@Test
public void testClusterResource() throws Exception {
    rmMock.availableRes(new MockResource(1024, 2));

    appMaster.setRmClient(rmMock);
    appMaster.setNmClient(new NMMock());

    props.cpusPerNode(8);
    props.memoryPerNode(10240);
    props.instances(3);

    Thread thread = runAppMaster(appMaster);

    List<AMRMClient.ContainerRequest> contRequests = collectRequests(rmMock, 1, 1000);

    interruptedThread(thread);

    assertEquals(0, contRequests.size());
}
 

@Override
public void startNimbus(int number, int containerMemory, int containerVirtualCores) throws TException {
    //set priority to 1 which identity this container is allocated for nimbus
    String dstHost = "*";
    if (!jstormMasterContext.previousNimbusHost.equals("") && jstormMasterContext.nimbusContainers.size() == 0)
        dstHost = jstormMasterContext.previousNimbusHost;

    LOG.info("dstHost:" + dstHost);
    AMRMClient.ContainerRequest containerAsk = jstormMaster.setupContainerAskForRM(containerMemory, containerVirtualCores, 1, dstHost);
    try {
        requestQueue.put(containerAsk);
    } catch (InterruptedException e) {
        e.printStackTrace();
    }
    amRMClient.addContainerRequest(containerAsk);
    jstormMasterContext.numRequestedContainers.getAndIncrement();
}
 

@Override
public synchronized void onNext(final ResourceRequestEvent resourceRequestEvent) {
  LOG.log(Level.FINEST, "Got ResourceRequestEvent in YarnResourceRequestHandler: memory = {0}, cores = {1}.",
      new Object[]{resourceRequestEvent.getMemorySize(), resourceRequestEvent.getVirtualCores()});

  final String[] nodes = resourceRequestEvent.getNodeNameList().size() == 0 ? null :
      resourceRequestEvent.getNodeNameList().toArray(new String[resourceRequestEvent.getNodeNameList().size()]);
  final String[] racks = resourceRequestEvent.getRackNameList().size() == 0 ? null :
      resourceRequestEvent.getRackNameList().toArray(new String[resourceRequestEvent.getRackNameList().size()]);

  // set the priority for the request
  final Priority pri = getPriority(resourceRequestEvent);
  final Resource resource = getResource(resourceRequestEvent);
  final boolean relaxLocality = resourceRequestEvent.getRelaxLocality().orElse(true);
  final String nodeLabelExpression = resourceRequestEvent.getNodeLabelExpression().orElse("");

  final AMRMClient.ContainerRequest[] containerRequests =
      new AMRMClient.ContainerRequest[resourceRequestEvent.getResourceCount()];

  for (int i = 0; i < resourceRequestEvent.getResourceCount(); i++) {
    containerRequests[i] =
        new AMRMClient.ContainerRequest(resource, nodes, racks, pri, relaxLocality,
            StringUtils.isEmpty(nodeLabelExpression) ? null : nodeLabelExpression);
  }
  this.yarnContainerRequestHandler.onContainerRequest(containerRequests);
}
 

/**
 * @param rmMock RM mock.
 * @param expectedCnt Expected cnt.
 * @param timeOut Timeout.
 * @return Requests.
 */
private List<AMRMClient.ContainerRequest> collectRequests(RMMock rmMock, int expectedCnt, int timeOut) {
    long startTime = System.currentTimeMillis();

    List<AMRMClient.ContainerRequest> requests = rmMock.requests();

    while (requests.size() < expectedCnt
       && (System.currentTimeMillis() - startTime) < timeOut)
        requests = rmMock.requests();

    return requests;
}
 

@SuppressWarnings("unchecked")
@Test(timeout=10000)
public void testMatchingRequestsForTopPriority() {
  String[] hosts = { "host1" };
  String[] racks = { "rack1" };
  AMRMClient.ContainerRequest req1 = new AMRMClient.ContainerRequest(
    Resource.newInstance(2048, 1), hosts, racks,
    Priority.newInstance(1));
  AMRMClient.ContainerRequest req2 = new AMRMClient.ContainerRequest(
    Resource.newInstance(1024, 1), hosts, racks,
    Priority.newInstance(2));
  AMRMClient.ContainerRequest req3 = new AMRMClient.ContainerRequest(
    Resource.newInstance(1024, 1), hosts, racks,
    Priority.newInstance(3));
  amrmClient.addContainerRequest(req1);
  amrmClient.addContainerRequest(req2);
  amrmClient.addContainerRequest(req3);

  Assert.assertTrue(amrmClient.getMatchingRequestsForTopPriority("host1",
    Resource.newInstance(1024, 1)).isEmpty());

  List<? extends Collection<AMRMClient.ContainerRequest>> ret =
    amrmClient.getMatchingRequestsForTopPriority("host1",
      Resource.newInstance(2048, 1));
  Assert.assertFalse(ret.isEmpty());
  Assert.assertEquals(req1, ret.get(0).iterator().next());

  amrmClient.removeContainerRequest(req1);

  ret = amrmClient.getMatchingRequestsForTopPriority("host1",
      Resource.newInstance(1024, 1));
  Assert.assertFalse(ret.isEmpty());
  Assert.assertEquals(req2, ret.get(0).iterator().next());
}
 

private AMRMClientAsync<AMRMClient.ContainerRequest> initClient(Configuration conf) throws YarnException, IOException {
    AMRMClientAsync.CallbackHandler allocListener = new AMRMClientAsync.CallbackHandler() {
        @Override
        public void onContainersCompleted(List<ContainerStatus> statuses) {
        }

        @Override
        public void onContainersAllocated(List<Container> containers) {
        }

        @Override
        public void onShutdownRequest() {
            LOG.warn("Shutting down");
            end.set(true);
        }

        @Override
        public void onNodesUpdated(List<NodeReport> updatedNodes) {
        }

        @Override
        public float getProgress() {
            return 0;
        }

        @Override
        public void onError(Throwable e) {
            LOG.error("Unexpected error", e);
            end.set(true);
        }
    };
    AMRMClientAsync<AMRMClient.ContainerRequest> rmClient = AMRMClientAsync.createAMRMClientAsync(1000, allocListener);
    rmClient.init(conf);
    rmClient.start();

    // Register with ResourceManager
    rmClient.registerApplicationMaster(Utils.localHostName(), 0, "");

    return rmClient;
}
 

public void requestContainers(int number, Resource characteristic) {
  Priority pri = Priority.newInstance(0);
  state.getQueue(characteristic);
  AMRMClient.ContainerRequest request = new AMRMClient.ContainerRequest(characteristic, null, null, pri, true);
  for(int i = 0;i < number;++i) {
    amRMClient.addContainerRequest(request);
  }
}
 

/**
 * 构建对于container的申请
 * @param containerType
 * @param mem
 * @return
 */
public static AMRMClient.ContainerRequest setupContainerAskForRM(int containerType, int mem) {
    Priority pri = Records.newRecord(Priority.class);
    pri.setPriority(containerType);
    Resource capability = Records.newRecord(Resource.class);
    capability.setMemory(mem);
    capability.setVirtualCores(1);
    AMRMClient.ContainerRequest request = new AMRMClient.ContainerRequest(capability, null, null, pri);
    LOG.error("#### 申请container:[" + request.toString()+"]");
    return request;
}
 

private boolean isSameKindOfRequest(final AMRMClient.ContainerRequest r1, final AMRMClient.ContainerRequest r2) {
  final boolean nodeLabelExpressionIsEqual = r1.getNodeLabelExpression() == r2.getNodeLabelExpression() ||
      (r1.getNodeLabelExpression() != null && r1.getNodeLabelExpression().equals(r2.getNodeLabelExpression()));
  return r1.getPriority().compareTo(r2.getPriority()) == 0
      && r1.getCapability().compareTo(r2.getCapability()) == 0
      && nodeLabelExpressionIsEqual
      && r1.getRelaxLocality() == r2.getRelaxLocality()
      && ListUtils.isEqualList(r1.getNodes(), r2.getNodes())
      && ListUtils.isEqualList(r1.getRacks(), r2.getRacks());
}
 

@Nonnull
@VisibleForTesting
AMRMClient.ContainerRequest getContainerRequest() {
	return new AMRMClient.ContainerRequest(
		getContainerResource(),
		null,
		null,
		RM_REQUEST_PRIORITY);
}
 

public void initialize(AMRMClientAsync<AMRMClient.ContainerRequest> amRMClient
                      , NMClientAsync nmClient
                      , ServiceDiscoverer serviceDiscoverer
                      )
{
  this.nmClient = nmClient;
  this.amRMClient = amRMClient;
  this.serviceDiscoverer = serviceDiscoverer;
}
 

protected AMRMClientAsync<AMRMClient.ContainerRequest> createAndStartResourceManagerClient(
		YarnConfiguration yarnConfiguration,
		int yarnHeartbeatIntervalMillis,
		@Nullable String webInterfaceUrl) throws Exception {
	AMRMClientAsync<AMRMClient.ContainerRequest> resourceManagerClient = AMRMClientAsync.createAMRMClientAsync(
		yarnHeartbeatIntervalMillis,
		this);

	resourceManagerClient.init(yarnConfiguration);
	resourceManagerClient.start();

	//TODO: change akka address to tcp host and port, the getAddress() interface should return a standard tcp address
	Tuple2<String, Integer> hostPort = parseHostPort(getAddress());

	final int restPort;

	if (webInterfaceUrl != null) {
		final int lastColon = webInterfaceUrl.lastIndexOf(':');

		if (lastColon == -1) {
			restPort = -1;
		} else {
			restPort = Integer.valueOf(webInterfaceUrl.substring(lastColon + 1));
		}
	} else {
		restPort = -1;
	}

	final RegisterApplicationMasterResponse registerApplicationMasterResponse =
		resourceManagerClient.registerApplicationMaster(hostPort.f0, restPort, webInterfaceUrl);
	getContainersFromPreviousAttempts(registerApplicationMasterResponse);

	return resourceManagerClient;
}
 

@Override
public void removeContainerRequest(AMRMClient.ContainerRequest req) {
	removeContainerRequestConsumer.accept(req, handler);
}
 

void setAddContainerRequestConsumer(
	BiConsumer<AMRMClient.ContainerRequest, CallbackHandler> addContainerRequestConsumer) {
	this.addContainerRequestConsumer = Preconditions.checkNotNull(addContainerRequestConsumer);
}
 

@Override
public List<? extends Collection<AMRMClient.ContainerRequest>> getMatchingRequests(Priority priority, String resourceName, Resource capability) {
	return getMatchingRequestsFunction.apply(Tuple4.of(priority, resourceName, capability, handler));
}
 

private void onContainersOfResourceAllocated(Resource resource, List<Container> containers) {
	final List<WorkerResourceSpec> pendingWorkerResourceSpecs =
		workerSpecContainerResourceAdapter.getWorkerSpecs(resource, matchingStrategy).stream()
			.flatMap(spec -> Collections.nCopies(getNumRequestedNotAllocatedWorkersFor(spec), spec).stream())
			.collect(Collectors.toList());

	int numPending = pendingWorkerResourceSpecs.size();
	log.info("Received {} containers with resource {}, {} pending container requests.",
		containers.size(),
		resource,
		numPending);

	final Iterator<Container> containerIterator = containers.iterator();
	final Iterator<WorkerResourceSpec> pendingWorkerSpecIterator = pendingWorkerResourceSpecs.iterator();
	final Iterator<AMRMClient.ContainerRequest> pendingRequestsIterator =
		getPendingRequestsAndCheckConsistency(resource, pendingWorkerResourceSpecs.size()).iterator();

	int numAccepted = 0;
	while (containerIterator.hasNext() && pendingWorkerSpecIterator.hasNext()) {
		final WorkerResourceSpec workerResourceSpec = pendingWorkerSpecIterator.next();
		final Container container = containerIterator.next();
		final AMRMClient.ContainerRequest pendingRequest = pendingRequestsIterator.next();
		final ResourceID resourceId = getContainerResourceId(container);

		notifyNewWorkerAllocated(workerResourceSpec, resourceId);
		startTaskExecutorInContainer(container, workerResourceSpec, resourceId);
		removeContainerRequest(pendingRequest, workerResourceSpec);

		numAccepted++;
	}
	numPending -= numAccepted;

	int numExcess = 0;
	while (containerIterator.hasNext()) {
		returnExcessContainer(containerIterator.next());
		numExcess++;
	}

	log.info("Accepted {} requested containers, returned {} excess containers, {} pending container requests of resource {}.",
		numAccepted, numExcess, numPending, resource);
}
 

AMRMClient.ContainerRequest[] getRequests() {
  return requests;
}
 

private void removeContainerRequest(AMRMClient.ContainerRequest pendingContainerRequest) {
	numPendingContainerRequests--;

	log.info("Removing container request {}. Pending container requests {}.", pendingContainerRequest, numPendingContainerRequests);

	resourceManagerClient.removeContainerRequest(pendingContainerRequest);
}
 

/**
 * Sets RMClient
 *
 * @param rmClient AMRMClientAsync.
 */
public void setRmClient(AMRMClientAsync<AMRMClient.ContainerRequest> rmClient) {
    this.rmClient = rmClient;
}