@Test
public void testTaskPullAtLaunching() {
WrappedContainer wc = new WrappedContainer();
wc.launchContainer();
AMContainerTask pulledTask = wc.pullTaskToRun();
wc.verifyState(AMContainerState.LAUNCHING);
wc.verifyNoOutgoingEvents();
assertFalse(pulledTask.shouldDie());
assertNull(pulledTask.getTask());
}
@Test
// Assign before launch.
public void tetSingleSuccessfulTaskFlow() {
WrappedContainer wc = new WrappedContainer();
wc.verifyState(AMContainerState.ALLOCATED);
// Launch request.
wc.launchContainer();
wc.verifyState(AMContainerState.LAUNCHING);
// 1 Launch request.
wc.verifyCountAndGetOutgoingEvents(1);
// Assign task.
wc.assignTaskAttempt(wc.taskAttemptID);
wc.verifyState(AMContainerState.LAUNCHING);
wc.verifyNoOutgoingEvents();
assertEquals(wc.taskAttemptID, wc.amContainer.getQueuedTaskAttempts().get(0));
// Container Launched
wc.containerLaunched();
wc.verifyState(AMContainerState.IDLE);
wc.verifyNoOutgoingEvents();
assertEquals(wc.taskAttemptID, wc.amContainer.getQueuedTaskAttempts().get(0));
assertNull(wc.amContainer.getRunningTaskAttempt());
verify(wc.tal).registerRunningContainer(wc.containerID);
verify(wc.chh).register(wc.containerID);
// Pull TA
AMContainerTask pulledTask = wc.pullTaskToRun();
wc.verifyState(AMContainerState.RUNNING);
wc.verifyNoOutgoingEvents();
assertFalse(pulledTask.shouldDie());
assertEquals(wc.taskSpec.getTaskAttemptID(), pulledTask.getTask().getTaskAttemptID());
assertEquals(wc.taskAttemptID, wc.amContainer.getRunningTaskAttempt());
assertEquals(0, wc.amContainer.getQueuedTaskAttempts().size());
wc.taskAttemptSucceeded(wc.taskAttemptID);
wc.verifyState(AMContainerState.IDLE);
wc.verifyNoOutgoingEvents();
assertNull(wc.amContainer.getRunningTaskAttempt());
verify(wc.tal).unregisterTaskAttempt(wc.taskAttemptID);
wc.containerCompleted(false);
wc.verifyHistoryStopEvent();
wc.verifyState(AMContainerState.COMPLETED);
wc.verifyNoOutgoingEvents();
verify(wc.tal).unregisterRunningContainer(wc.containerID);
verify(wc.chh).unregister(wc.containerID);
assertEquals(1, wc.amContainer.getAllTaskAttempts().size());
assertFalse(wc.amContainer.isInErrorState());
}
@Test
public void testLocalResourceAddition() {
WrappedContainer wc = new WrappedContainer();
String rsrc1 = "rsrc1";
String rsrc2 = "rsrc2";
String rsrc3 = "rsrc3";
Map<String, LocalResource> initialResources = Maps.newHashMap();
initialResources.put(rsrc1, createLocalResource(rsrc1));
wc.launchContainer(initialResources, new Credentials());
wc.containerLaunched();
wc.assignTaskAttempt(wc.taskAttemptID);
AMContainerTask task1 = wc.pullTaskToRun();
assertEquals(0, task1.getAdditionalResources().size());
wc.taskAttemptSucceeded(wc.taskAttemptID);
// Add some resources to the next task.
Map<String, LocalResource> additionalResources = Maps.newHashMap();
additionalResources.put(rsrc2, createLocalResource(rsrc2));
additionalResources.put(rsrc3, createLocalResource(rsrc3));
TezTaskAttemptID taID2 = TezTaskAttemptID.getInstance(wc.taskID, 2);
wc.assignTaskAttempt(taID2, additionalResources, new Credentials());
AMContainerTask task2 = wc.pullTaskToRun();
Map<String, LocalResource> pullTaskAdditionalResources = task2.getAdditionalResources();
assertEquals(2, pullTaskAdditionalResources.size());
pullTaskAdditionalResources.remove(rsrc2);
pullTaskAdditionalResources.remove(rsrc3);
assertEquals(0, pullTaskAdditionalResources.size());
wc.taskAttemptSucceeded(taID2);
// Verify Resources registered for this container.
Map<String, LocalResource> containerLRs =
new HashMap<String, LocalResource>(wc.amContainer.containerLocalResources);
assertEquals(3, containerLRs.size());
containerLRs.remove(rsrc1);
containerLRs.remove(rsrc2);
containerLRs.remove(rsrc3);
assertEquals(0, containerLRs.size());
// Try launching another task with the same reosurces as Task2. Verify the
// task is not asked to re-localize again.
TezTaskAttemptID taID3 = TezTaskAttemptID.getInstance(wc.taskID, 3);
wc.assignTaskAttempt(taID3, new HashMap<String, LocalResource>(), new Credentials());
AMContainerTask task3 = wc.pullTaskToRun();
assertEquals(0, task3.getAdditionalResources().size());
wc.taskAttemptSucceeded(taID3);
// Verify references are cleared after a container completes.
wc.containerCompleted(false);
assertNull(wc.amContainer.containerLocalResources);
assertNull(wc.amContainer.additionalLocalResources);
}
@SuppressWarnings("unchecked")
@Test
public void testCredentialsTransfer() {
WrappedContainerMultipleDAGs wc = new WrappedContainerMultipleDAGs();
TezDAGID dagID2 = TezDAGID.getInstance("800", 500, 2);
TezDAGID dagID3 = TezDAGID.getInstance("800", 500, 3);
TezVertexID vertexID2 = TezVertexID.getInstance(dagID2, 1);
TezVertexID vertexID3 = TezVertexID.getInstance(dagID3, 1);
TezTaskID taskID2 = TezTaskID.getInstance(vertexID2, 1);
TezTaskID taskID3 = TezTaskID.getInstance(vertexID3, 1);
TezTaskAttemptID attempt11 = TezTaskAttemptID.getInstance(wc.taskID, 200);
TezTaskAttemptID attempt12 = TezTaskAttemptID.getInstance(wc.taskID, 300);
TezTaskAttemptID attempt21 = TezTaskAttemptID.getInstance(taskID2, 200);
TezTaskAttemptID attempt22 = TezTaskAttemptID.getInstance(taskID2, 300);
TezTaskAttemptID attempt31 = TezTaskAttemptID.getInstance(taskID3, 200);
TezTaskAttemptID attempt32 = TezTaskAttemptID.getInstance(taskID3, 300);
Map<String, LocalResource> LRs = new HashMap<String, LocalResource>();
AMContainerTask fetchedTask = null;
Token<TokenIdentifier> amGenToken = mock(Token.class);
Token<TokenIdentifier> token1 = mock(Token.class);
Token<TokenIdentifier> token3 = mock(Token.class);
Credentials containerCredentials = new Credentials();
TokenCache.setSessionToken(amGenToken, containerCredentials);
Text token1Name = new Text("tokenDag1");
Text token3Name = new Text("tokenDag3");
Credentials dag1Credentials = new Credentials();
dag1Credentials.addToken(new Text(token1Name), token1);
Credentials dag3Credentials = new Credentials();
dag3Credentials.addToken(new Text(token3Name), token3);
wc.launchContainer(new HashMap<String, LocalResource>(), containerCredentials);
wc.containerLaunched();
wc.assignTaskAttempt(attempt11, LRs, dag1Credentials);
fetchedTask = wc.pullTaskToRun();
assertTrue(fetchedTask.haveCredentialsChanged());
assertNotNull(fetchedTask.getCredentials());
assertNotNull(fetchedTask.getCredentials().getToken(token1Name));
wc.taskAttemptSucceeded(attempt11);
wc.assignTaskAttempt(attempt12, LRs, dag1Credentials);
fetchedTask = wc.pullTaskToRun();
assertFalse(fetchedTask.haveCredentialsChanged());
assertNull(fetchedTask.getCredentials());
wc.taskAttemptSucceeded(attempt12);
// Move to running a second DAG, with no credentials.
wc.setNewDAGID(dagID2);
wc.assignTaskAttempt(attempt21, LRs, null);
fetchedTask = wc.pullTaskToRun();
assertTrue(fetchedTask.haveCredentialsChanged());
assertNull(fetchedTask.getCredentials());
wc.taskAttemptSucceeded(attempt21);
wc.assignTaskAttempt(attempt22, LRs, null);
fetchedTask = wc.pullTaskToRun();
assertFalse(fetchedTask.haveCredentialsChanged());
assertNull(fetchedTask.getCredentials());
wc.taskAttemptSucceeded(attempt22);
// Move to running a third DAG, with Credentials this time
wc.setNewDAGID(dagID3);
wc.assignTaskAttempt(attempt31, LRs, dag3Credentials);
fetchedTask = wc.pullTaskToRun();
assertTrue(fetchedTask.haveCredentialsChanged());
assertNotNull(fetchedTask.getCredentials());
assertNotNull(fetchedTask.getCredentials().getToken(token3Name));
assertNull(fetchedTask.getCredentials().getToken(token1Name));
wc.taskAttemptSucceeded(attempt31);
wc.assignTaskAttempt(attempt32, LRs, dag1Credentials);
fetchedTask = wc.pullTaskToRun();
assertFalse(fetchedTask.haveCredentialsChanged());
assertNull(fetchedTask.getCredentials());
wc.taskAttemptSucceeded(attempt32);
}