Example #1
0
 public void launchContainer(
     Map<String, LocalResource> localResources, Credentials credentials) {
   reset(eventHandler);
   @SuppressWarnings("unchecked")
   Token<JobTokenIdentifier> jobToken = mock(Token.class);
   TokenCache.setSessionToken(jobToken, credentials);
   amContainer.handle(
       new AMContainerEventLaunchRequest(
           containerID,
           vertexID,
           new ContainerContext(
               localResources, credentials, new HashMap<String, String>(), "")));
 }
Example #2
0
 @VisibleForTesting
 public TaskRunnerCallable(
     SubmitWorkRequestProto request,
     QueryFragmentInfo fragmentInfo,
     Configuration conf,
     ExecutionContext executionContext,
     Map<String, String> envMap,
     Credentials credentials,
     long memoryAvailable,
     AMReporter amReporter,
     ConfParams confParams,
     LlapDaemonExecutorMetrics metrics,
     KilledTaskHandler killedTaskHandler,
     FragmentCompletionHandler fragmentCompleteHandler) {
   this.request = request;
   this.fragmentInfo = fragmentInfo;
   this.conf = conf;
   this.executionContext = executionContext;
   this.envMap = envMap;
   this.objectRegistry = new ObjectRegistryImpl();
   this.credentials = credentials;
   this.memoryAvailable = memoryAvailable;
   this.confParams = confParams;
   this.jobToken = TokenCache.getSessionToken(credentials);
   this.taskSpec = Converters.getTaskSpecfromProto(request.getFragmentSpec());
   this.amReporter = amReporter;
   // Register with the AMReporter when the callable is setup. Unregister once it starts running.
   if (jobToken != null) {
     this.amReporter.registerTask(
         request.getAmHost(),
         request.getAmPort(),
         request.getUser(),
         jobToken,
         null,
         request.getFragmentSpec().getDagName());
   }
   this.metrics = metrics;
   this.requestId = request.getFragmentSpec().getFragmentIdentifierString();
   // TODO Change this to the queryId/Name when that's available.
   this.queryId = request.getFragmentSpec().getDagName();
   this.killedTaskHandler = killedTaskHandler;
   this.fragmentCompletionHanler = fragmentCompleteHandler;
 }
Example #3
0
  @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);
  }