Java Code Examples for org.apache.hadoop.yarn.client.api.async.AMRMClientAsync#registerApplicationMaster()

@Test (timeout = 10000)
public void testAMRMClientAsyncShutDown() throws Exception {
  Configuration conf = new Configuration();
  TestCallbackHandler callbackHandler = new TestCallbackHandler();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);

  createAllocateResponse(new ArrayList<ContainerStatus>(),
    new ArrayList<Container>(), null);
  when(client.allocate(anyFloat())).thenThrow(
    new ApplicationAttemptNotFoundException("app not found, shut down"));

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

  asyncClient.registerApplicationMaster("localhost", 1234, null);

  Thread.sleep(50);

  verify(client, times(1)).allocate(anyFloat());
  asyncClient.stop();
}
 

@Test (timeout = 10000)
public void testAMRMClientAsyncShutDown() throws Exception {
  Configuration conf = new Configuration();
  TestCallbackHandler callbackHandler = new TestCallbackHandler();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);

  createAllocateResponse(new ArrayList<ContainerStatus>(),
    new ArrayList<Container>(), null);
  when(client.allocate(anyFloat())).thenThrow(
    new ApplicationAttemptNotFoundException("app not found, shut down"));

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

  asyncClient.registerApplicationMaster("localhost", 1234, null);

  Thread.sleep(50);

  verify(client, times(1)).allocate(anyFloat());
  asyncClient.stop();
}
 

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

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

public void run() throws Exception {
  int virtualCores = Integer.parseInt(cli.getOptionValue("virtualCores", "1"));

  AMRMClientAsync<ContainerRequest> rmClient = AMRMClientAsync.createAMRMClientAsync(100, this);
  rmClient.init(getConfiguration());
  rmClient.start();

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

  // Priority for worker containers - priorities are intra-application
  Priority priority = Records.newRecord(Priority.class);
  priority.setPriority(0);

  // Resource requirements for worker containers
  Resource capability = Records.newRecord(Resource.class);
  capability.setMemory(memory);
  capability.setVirtualCores(virtualCores);

  // Make container requests to ResourceManager
  for (int i = 0; i < numContainersToWaitFor; ++i)
    rmClient.addContainerRequest(new ContainerRequest(capability, null, null, priority));

  log.info("Waiting for " + numContainersToWaitFor + " containers to finish");
  while (!doneWithContainers())
    Thread.sleep(10000);

  log.info("SolrMaster application shutdown.");

  // Un-register with ResourceManager
  try {
    rmClient.unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, "", "");
  } catch (Exception exc) {
    // safe to ignore ... this usually fails anyway
  }
}
 

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);
	updateMatchingStrategy(registerApplicationMasterResponse);

	return resourceManagerClient;
}
 

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

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

@Test (timeout = 10000)
public void testAMRMClientAsyncShutDownWithWaitFor() throws Exception {
  Configuration conf = new Configuration();
  final TestCallbackHandler callbackHandler = new TestCallbackHandler();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);
  when(client.allocate(anyFloat())).thenThrow(
    new ApplicationAttemptNotFoundException("app not found, shut down"));

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

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

  asyncClient.registerApplicationMaster("localhost", 1234, null);
  asyncClient.waitFor(checker);

  asyncClient.stop();
  // stopping should have joined all threads and completed all callbacks
  Assert.assertTrue(callbackHandler.callbackCount == 0);

  verify(client, times(1)).allocate(anyFloat());
  asyncClient.stop();
}
 

private void runHeartBeatThrowOutException(Exception ex) throws Exception{
  Configuration conf = new Configuration();
  TestCallbackHandler callbackHandler = new TestCallbackHandler();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);
  when(client.allocate(anyFloat())).thenThrow(ex);

  AMRMClientAsync<ContainerRequest> asyncClient =
      AMRMClientAsync.createAMRMClientAsync(client, 20, callbackHandler);
  asyncClient.init(conf);
  asyncClient.start();
  
  synchronized (callbackHandler.notifier) {
    asyncClient.registerApplicationMaster("localhost", 1234, null);
    while(callbackHandler.savedException == null) {
      try {
        callbackHandler.notifier.wait();
      } catch (InterruptedException e) {
        e.printStackTrace();
      }
    }
  }
  Assert.assertTrue(callbackHandler.savedException.getMessage().contains(
      ex.getMessage()));
  
  asyncClient.stop();
  // stopping should have joined all threads and completed all callbacks
  Assert.assertTrue(callbackHandler.callbackCount == 0);
}
 

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

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

@Test (timeout = 10000)
public void testAMRMClientAsyncShutDownWithWaitFor() throws Exception {
  Configuration conf = new Configuration();
  final TestCallbackHandler callbackHandler = new TestCallbackHandler();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);
  when(client.allocate(anyFloat())).thenThrow(
    new ApplicationAttemptNotFoundException("app not found, shut down"));

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

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

  asyncClient.registerApplicationMaster("localhost", 1234, null);
  asyncClient.waitFor(checker);

  asyncClient.stop();
  // stopping should have joined all threads and completed all callbacks
  Assert.assertTrue(callbackHandler.callbackCount == 0);

  verify(client, times(1)).allocate(anyFloat());
  asyncClient.stop();
}
 

private void runHeartBeatThrowOutException(Exception ex) throws Exception{
  Configuration conf = new Configuration();
  TestCallbackHandler callbackHandler = new TestCallbackHandler();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);
  when(client.allocate(anyFloat())).thenThrow(ex);

  AMRMClientAsync<ContainerRequest> asyncClient =
      AMRMClientAsync.createAMRMClientAsync(client, 20, callbackHandler);
  asyncClient.init(conf);
  asyncClient.start();
  
  synchronized (callbackHandler.notifier) {
    asyncClient.registerApplicationMaster("localhost", 1234, null);
    while(callbackHandler.savedException == null) {
      try {
        callbackHandler.notifier.wait();
      } catch (InterruptedException e) {
        e.printStackTrace();
      }
    }
  }
  Assert.assertTrue(callbackHandler.savedException.getMessage().contains(
      ex.getMessage()));
  
  asyncClient.stop();
  // stopping should have joined all threads and completed all callbacks
  Assert.assertTrue(callbackHandler.callbackCount == 0);
}
 

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

void runCallBackThrowOutException(TestCallbackHandler2 callbackHandler) throws
      InterruptedException, YarnException, IOException {
  Configuration conf = new Configuration();
  @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;
  callbackHandler.throwOutException = true;
  asyncClient.init(conf);
  asyncClient.start();

  // call register and wait for error callback and stop
  synchronized (callbackHandler.notifier) {
    asyncClient.registerApplicationMaster("localhost", 1234, null);
    while(callbackHandler.notify == false) {
      try {
        callbackHandler.notifier.wait();
      } catch (InterruptedException e) {
        e.printStackTrace();
      }
    }
  }
  // verify error invoked
  verify(callbackHandler, times(0)).getProgress();
  verify(callbackHandler, times(1)).onError(any(Exception.class));
  // sleep to wait for a few heartbeat calls that can trigger callbacks
  Thread.sleep(50);
  // verify no more invocations after the first one.
  // ie. callback thread has stopped
  verify(callbackHandler, times(0)).getProgress();
  verify(callbackHandler, times(1)).onError(any(Exception.class));
}
 

void runCallBackThrowOutException(TestCallbackHandler2 callbackHandler) throws
      InterruptedException, YarnException, IOException {
  Configuration conf = new Configuration();
  @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;
  callbackHandler.throwOutException = true;
  asyncClient.init(conf);
  asyncClient.start();

  // call register and wait for error callback and stop
  synchronized (callbackHandler.notifier) {
    asyncClient.registerApplicationMaster("localhost", 1234, null);
    while(callbackHandler.notify == false) {
      try {
        callbackHandler.notifier.wait();
      } catch (InterruptedException e) {
        e.printStackTrace();
      }
    }
  }
  // verify error invoked
  verify(callbackHandler, times(0)).getProgress();
  verify(callbackHandler, times(1)).onError(any(Exception.class));
  // sleep to wait for a few heartbeat calls that can trigger callbacks
  Thread.sleep(50);
  // verify no more invocations after the first one.
  // ie. callback thread has stopped
  verify(callbackHandler, times(0)).getProgress();
  verify(callbackHandler, times(1)).onError(any(Exception.class));
}
 

@Test
public void testAmShutdown() throws IOException, YarnException {

  Assume.assumeTrue(
      "This test requires a YARN Resource Manager to connect to",
      Boolean.parseBoolean(System.getenv("REEF_TEST_YARN")));

  final YarnConfiguration yarnConfig = new YarnConfiguration();

  // Start YARN client and register the application

  final YarnClient yarnClient = YarnClient.createYarnClient();
  yarnClient.init(yarnConfig);
  yarnClient.start();

  final ContainerLaunchContext containerContext = Records.newRecord(ContainerLaunchContext.class);
  containerContext.setCommands(Collections.<String>emptyList());
  containerContext.setLocalResources(Collections.<String, LocalResource>emptyMap());
  containerContext.setEnvironment(Collections.<String, String>emptyMap());
  containerContext.setTokens(getTokens());

  final ApplicationSubmissionContext appContext = yarnClient.createApplication().getApplicationSubmissionContext();
  appContext.setApplicationName("REEF_Unmanaged_AM_Test");
  appContext.setAMContainerSpec(containerContext);
  appContext.setUnmanagedAM(true);
  appContext.setQueue("default");

  final ApplicationId applicationId = appContext.getApplicationId();
  LOG.log(Level.INFO, "Registered YARN application: {0}", applicationId);

  yarnClient.submitApplication(appContext);

  LOG.log(Level.INFO, "YARN application submitted: {0}", applicationId);

  addToken(yarnClient.getAMRMToken(applicationId));

  // Start the AM

  final AMRMClientAsync<AMRMClient.ContainerRequest> rmClient = AMRMClientAsync.createAMRMClientAsync(1000, this);
  rmClient.init(yarnConfig);
  rmClient.start();

  final NMClientAsync nmClient = new NMClientAsyncImpl(this);
  nmClient.init(yarnConfig);
  nmClient.start();

  final RegisterApplicationMasterResponse registration =
      rmClient.registerApplicationMaster(NetUtils.getHostname(), -1, null);

  LOG.log(Level.INFO, "Unmanaged AM is running: {0}", registration);

  rmClient.unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, "Success!", null);

  LOG.log(Level.INFO, "Unregistering AM: state {0}", rmClient.getServiceState());

  // Shutdown the AM

  rmClient.stop();
  nmClient.stop();

  // Get the final application report

  final ApplicationReport appReport = yarnClient.getApplicationReport(applicationId);
  final YarnApplicationState appState = appReport.getYarnApplicationState();
  final FinalApplicationStatus finalAttemptStatus = appReport.getFinalApplicationStatus();

  LOG.log(Level.INFO, "Application {0} final attempt {1} status: {2}/{3}", new Object[] {
      applicationId, appReport.getCurrentApplicationAttemptId(), appState, finalAttemptStatus});

  Assert.assertEquals("Application must be in FINISHED state", YarnApplicationState.FINISHED, appState);
  Assert.assertEquals("Final status must be SUCCEEDED", FinalApplicationStatus.SUCCEEDED, finalAttemptStatus);

  // Shutdown YARN client

  yarnClient.stop();
}