Example usage for org.apache.hadoop.yarn.api.protocolrecords RegisterApplicationMasterResponse getApplicationACLs

List of usage examples for org.apache.hadoop.yarn.api.protocolrecords RegisterApplicationMasterResponse getApplicationACLs

Introduction

In this page you can find the example usage for org.apache.hadoop.yarn.api.protocolrecords RegisterApplicationMasterResponse getApplicationACLs.

Prototype

@Public
@Stable
public abstract Map<ApplicationAccessType, String> getApplicationACLs();

Source Link

Document

Get the ApplicationACLs for the application.

Usage

From source file:org.apache.hoya.yarn.appmaster.HoyaAppMaster.java

License:Apache License

/**
 * Create and run the cluster./*  w w  w.ja  v a2  s  .  co m*/
 * @return exit code
 * @throws Throwable on a failure
 */
private int createAndRunCluster(String clustername) throws Throwable {
    HoyaVersionInfo.loadAndPrintVersionInfo(log);

    //load the cluster description from the cd argument
    String hoyaClusterDir = serviceArgs.getHoyaClusterURI();
    URI hoyaClusterURI = new URI(hoyaClusterDir);
    Path clusterDirPath = new Path(hoyaClusterURI);
    HoyaFileSystem fs = getClusterFS();

    // build up information about the running application -this
    // will be passed down to the cluster status
    MapOperations appInformation = new MapOperations();

    AggregateConf instanceDefinition = InstanceIO.loadInstanceDefinitionUnresolved(fs, clusterDirPath);

    log.info("Deploying cluster {}:", instanceDefinition);

    //REVISIT: why is this done?
    appState.updateInstanceDefinition(instanceDefinition);
    File confDir = getLocalConfDir();
    if (!confDir.exists() || !confDir.isDirectory()) {
        log.error("Bad conf dir {}", confDir);
        File parentFile = confDir.getParentFile();
        log.error("Parent dir {}:\n{}", parentFile, HoyaUtils.listDir(parentFile));
        throw new BadCommandArgumentsException("Configuration directory %s doesn't exist", confDir);
    }

    Configuration serviceConf = getConfig();
    // Try to get the proper filtering of static resources through the yarn proxy working
    serviceConf.set("hadoop.http.filter.initializers",
            "org.apache.hadoop.yarn.server.webproxy.amfilter.AmFilterInitializer");

    conf = new YarnConfiguration(serviceConf);
    //get our provider
    MapOperations globalOptions = instanceDefinition.getInternalOperations().getGlobalOptions();
    String providerType = globalOptions.getMandatoryOption(OptionKeys.INTERNAL_PROVIDER_NAME);
    log.info("Cluster provider type is {}", providerType);
    HoyaProviderFactory factory = HoyaProviderFactory.createHoyaProviderFactory(providerType);
    providerService = factory.createServerProvider();
    // init the provider BUT DO NOT START IT YET
    providerService.init(getConfig());
    addService(providerService);

    InetSocketAddress address = HoyaUtils.getRmSchedulerAddress(conf);
    log.info("RM is at {}", address);
    yarnRPC = YarnRPC.create(conf);

    /*
     * Extract the container ID. This is then
     * turned into an (incompete) container
     */
    appMasterContainerID = ConverterUtils.toContainerId(
            HoyaUtils.mandatoryEnvVariable(ApplicationConstants.Environment.CONTAINER_ID.name()));
    appAttemptID = appMasterContainerID.getApplicationAttemptId();

    ApplicationId appid = appAttemptID.getApplicationId();
    log.info("Hoya AM for ID {}", appid.getId());

    appInformation.put(StatusKeys.INFO_AM_CONTAINER_ID, appMasterContainerID.toString());
    appInformation.put(StatusKeys.INFO_AM_APP_ID, appid.toString());
    appInformation.put(StatusKeys.INFO_AM_ATTEMPT_ID, appAttemptID.toString());

    UserGroupInformation currentUser = UserGroupInformation.getCurrentUser();
    Credentials credentials = currentUser.getCredentials();
    DataOutputBuffer dob = new DataOutputBuffer();
    credentials.writeTokenStorageToStream(dob);
    dob.close();
    // Now remove the AM->RM token so that containers cannot access it.
    Iterator<Token<?>> iter = credentials.getAllTokens().iterator();
    while (iter.hasNext()) {
        Token<?> token = iter.next();
        log.info("Token {}", token.getKind());
        if (token.getKind().equals(AMRMTokenIdentifier.KIND_NAME)) {
            iter.remove();
        }
    }
    allTokens = ByteBuffer.wrap(dob.getData(), 0, dob.getLength());

    // set up secret manager
    secretManager = new ClientToAMTokenSecretManager(appAttemptID, null);

    // if not a secure cluster, extract the username -it will be
    // propagated to workers
    if (!UserGroupInformation.isSecurityEnabled()) {
        hoyaUsername = System.getenv(HADOOP_USER_NAME);
        log.info(HADOOP_USER_NAME + "='{}'", hoyaUsername);
    }

    Map<String, String> envVars;

    /**
     * It is critical this section is synchronized, to stop async AM events
     * arriving while registering a restarting AM.
     */
    synchronized (appState) {
        int heartbeatInterval = HEARTBEAT_INTERVAL;

        //add the RM client -this brings the callbacks in
        asyncRMClient = AMRMClientAsync.createAMRMClientAsync(heartbeatInterval, this);
        addService(asyncRMClient);
        //wrap it for the app state model
        rmOperationHandler = new AsyncRMOperationHandler(asyncRMClient);
        //now bring it up
        runChildService(asyncRMClient);

        //nmclient relays callbacks back to this class
        nmClientAsync = new NMClientAsyncImpl("nmclient", this);
        runChildService(nmClientAsync);

        //bring up the Hoya RPC service
        startHoyaRPCServer();

        InetSocketAddress rpcServiceAddr = rpcService.getConnectAddress();
        appMasterHostname = rpcServiceAddr.getHostName();
        appMasterRpcPort = rpcServiceAddr.getPort();
        appMasterTrackingUrl = null;
        log.info("AM Server is listening at {}:{}", appMasterHostname, appMasterRpcPort);
        appInformation.put(StatusKeys.INFO_AM_HOSTNAME, appMasterHostname);
        appInformation.set(StatusKeys.INFO_AM_RPC_PORT, appMasterRpcPort);

        //build the role map
        List<ProviderRole> providerRoles = new ArrayList<ProviderRole>(providerService.getRoles());
        providerRoles.addAll(HoyaAMClientProvider.ROLES);

        // Start up the WebApp and track the URL for it
        webApp = new HoyaAMWebApp();
        WebApps.$for("hoyaam", WebAppApi.class, new WebAppApiImpl(this, appState, providerService), "ws")
                .with(serviceConf).start(webApp);
        appMasterTrackingUrl = "http://" + appMasterHostname + ":" + webApp.port();
        WebAppService<HoyaAMWebApp> webAppService = new WebAppService<HoyaAMWebApp>("hoya", webApp);

        webAppService.init(conf);
        webAppService.start();
        addService(webAppService);

        appInformation.put(StatusKeys.INFO_AM_WEB_URL, appMasterTrackingUrl + "/");
        appInformation.set(StatusKeys.INFO_AM_WEB_PORT, webApp.port());

        // Register self with ResourceManager
        // This will start heartbeating to the RM
        // address = HoyaUtils.getRmSchedulerAddress(asyncRMClient.getConfig());
        log.info("Connecting to RM at {},address tracking URL={}", appMasterRpcPort, appMasterTrackingUrl);
        RegisterApplicationMasterResponse response = asyncRMClient.registerApplicationMaster(appMasterHostname,
                appMasterRpcPort, appMasterTrackingUrl);
        Resource maxResources = response.getMaximumResourceCapability();
        containerMaxMemory = maxResources.getMemory();
        containerMaxCores = maxResources.getVirtualCores();
        appState.setContainerLimits(maxResources.getMemory(), maxResources.getVirtualCores());
        // set the RM-defined maximum cluster values
        appInformation.put(ResourceKeys.YARN_CORES, Integer.toString(containerMaxCores));
        appInformation.put(ResourceKeys.YARN_MEMORY, Integer.toString(containerMaxMemory));

        boolean securityEnabled = UserGroupInformation.isSecurityEnabled();
        if (securityEnabled) {
            secretManager.setMasterKey(response.getClientToAMTokenMasterKey().array());
            applicationACLs = response.getApplicationACLs();

            //tell the server what the ACLs are 
            rpcService.getServer().refreshServiceAcl(conf, new HoyaAMPolicyProvider());
        }

        // extract container list
        List<Container> liveContainers = AMRestartSupport.retrieveContainersFromPreviousAttempt(response);
        String amRestartSupported = Boolean.toString(liveContainers != null);
        appInformation.put(StatusKeys.INFO_AM_RESTART_SUPPORTED, amRestartSupported);

        //now validate the installation
        Configuration providerConf = providerService.loadProviderConfigurationInformation(confDir);

        providerService.validateApplicationConfiguration(instanceDefinition, confDir, securityEnabled);

        //determine the location for the role history data
        Path historyDir = new Path(clusterDirPath, HISTORY_DIR_NAME);

        //build the instance
        appState.buildInstance(instanceDefinition, providerConf, providerRoles, fs.getFileSystem(), historyDir,
                liveContainers, appInformation);

        // add the AM to the list of nodes in the cluster

        appState.buildAppMasterNode(appMasterContainerID, appMasterHostname, webApp.port(),
                appMasterHostname + ":" + webApp.port());

        // build up environment variables that the AM wants set in every container
        // irrespective of provider and role.
        envVars = new HashMap<String, String>();
        if (hoyaUsername != null) {
            envVars.put(HADOOP_USER_NAME, hoyaUsername);
        }
    }
    String rolesTmpSubdir = appMasterContainerID.toString() + "/roles";

    String amTmpDir = globalOptions.getMandatoryOption(OptionKeys.INTERNAL_AM_TMP_DIR);

    Path tmpDirPath = new Path(amTmpDir);
    Path launcherTmpDirPath = new Path(tmpDirPath, rolesTmpSubdir);
    fs.getFileSystem().mkdirs(launcherTmpDirPath);

    //launcher service
    launchService = new RoleLaunchService(this, providerService, fs, new Path(getGeneratedConfDir()), envVars,
            launcherTmpDirPath);

    runChildService(launchService);

    appState.noteAMLaunched();

    //Give the provider restricted access to the state
    providerService.bind(appState);

    // launch the provider; this is expected to trigger a callback that
    // brings up the service
    launchProviderService(instanceDefinition, confDir);

    try {
        //now block waiting to be told to exit the process
        waitForAMCompletionSignal();
        //shutdown time
    } finally {
        finish();
    }

    return amExitCode;
}

From source file:org.apache.tez.dag.app.rm.DagAwareYarnTaskScheduler.java

License:Apache License

@Override
public void start() throws Exception {
    super.start();
    client.start();//from www .  j  ava2s.  c o  m
    if (shouldReuseContainers) {
        reuseExecutor = createExecutor();
    }
    TaskSchedulerContext ctx = getContext();
    RegisterApplicationMasterResponse response = client.registerApplicationMaster(ctx.getAppHostName(),
            ctx.getAppClientPort(), ctx.getAppTrackingUrl());
    ctx.setApplicationRegistrationData(response.getMaximumResourceCapability(), response.getApplicationACLs(),
            response.getClientToAMTokenMasterKey(), response.getQueue());
    if (response.getSchedulerResourceTypes().contains(SchedulerResourceTypes.CPU)) {
        resourceCalculator = new MemCpuResourceCalculator();
    } else {
        resourceCalculator = new MemResourceCalculator();
    }
}

From source file:org.apache.tez.dag.app.rm.TaskScheduler.java

License:Apache License

@Override
public void serviceStart() {
    try {//from  w w w. j  ava2  s.co m
        RegisterApplicationMasterResponse response = null;
        synchronized (this) {
            amRmClient.start();
            response = amRmClient.registerApplicationMaster(appHostName, appHostPort, appTrackingUrl);
        }
        // upcall to app outside locks
        appClient.setApplicationRegistrationData(response.getMaximumResourceCapability(),
                response.getApplicationACLs());
    } catch (YarnException e) {
        LOG.error("Yarn Exception while registering", e);
        throw new TezUncheckedException(e);
    } catch (IOException e) {
        LOG.error("IO Exception while registering", e);
        throw new TezUncheckedException(e);
    }
}

From source file:org.apache.tez.dag.app.rm.TestDagAwareYarnTaskScheduler.java

License:Apache License

@SuppressWarnings({ "unchecked" })
@Test(timeout = 30000)/*from  www .  j av  a  2s.  c o m*/
public void testNoReuse() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock);

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

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
            regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
            regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Object mockTask1 = new MockTask("task1");
    Object mockCookie1 = new Object();
    Resource mockCapability = Resources.createResource(1024, 1);
    String[] hosts = { "host1", "host5" };
    String[] racks = { "/default-rack", "/default-rack" };
    Priority mockPriority = Priority.newInstance(1);
    ArgumentCaptor<TaskRequest> requestCaptor = ArgumentCaptor.forClass(TaskRequest.class);
    // allocate task
    scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1);
    drainableAppCallback.drain();
    verify(mockRMClient, times(1)).addContainerRequest(any(TaskRequest.class));

    // returned from task requests before allocation happens
    assertFalse(scheduler.deallocateTask(mockTask1, true, null, null));
    verify(mockApp, times(0)).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, times(1)).removeContainerRequest(any(TaskRequest.class));
    verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());

    // deallocating unknown task
    assertFalse(scheduler.deallocateTask(mockTask1, true, null, null));
    verify(mockApp, times(0)).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, times(1)).removeContainerRequest(any(TaskRequest.class));
    verify(mockRMClient, times(0)).releaseAssignedContainer((ContainerId) any());

    // allocate tasks
    Object mockTask2 = new MockTask("task2");
    Object mockCookie2 = new Object();
    Object mockTask3 = new MockTask("task3");
    Object mockCookie3 = new Object();
    scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1);
    drainableAppCallback.drain();
    verify(mockRMClient, times(2)).addContainerRequest(requestCaptor.capture());
    TaskRequest request1 = requestCaptor.getValue();
    scheduler.allocateTask(mockTask2, mockCapability, hosts, racks, mockPriority, null, mockCookie2);
    drainableAppCallback.drain();
    verify(mockRMClient, times(3)).addContainerRequest(requestCaptor.capture());
    TaskRequest request2 = requestCaptor.getValue();
    scheduler.allocateTask(mockTask3, mockCapability, hosts, racks, mockPriority, null, mockCookie3);
    drainableAppCallback.drain();
    verify(mockRMClient, times(4)).addContainerRequest(requestCaptor.capture());
    TaskRequest request3 = requestCaptor.getValue();

    NodeId host1 = NodeId.newInstance("host1", 1);
    NodeId host2 = NodeId.newInstance("host2", 2);
    NodeId host3 = NodeId.newInstance("host3", 3);
    NodeId host4 = NodeId.newInstance("host4", 4);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId mockCId1 = ContainerId.newContainerId(attemptId, 1);
    Container mockContainer1 = Container.newInstance(mockCId1, host1, null, mockCapability, mockPriority, null);
    ContainerId mockCId2 = ContainerId.newContainerId(attemptId, 2);
    Container mockContainer2 = Container.newInstance(mockCId2, host2, null, mockCapability, mockPriority, null);
    ContainerId mockCId3 = ContainerId.newContainerId(attemptId, 3);
    Container mockContainer3 = Container.newInstance(mockCId3, host3, null, mockCapability, mockPriority, null);
    ContainerId mockCId4 = ContainerId.newContainerId(attemptId, 4);
    Container mockContainer4 = Container.newInstance(mockCId4, host4, null, mockCapability, mockPriority, null);
    List<Container> containers = new ArrayList<>();
    containers.add(mockContainer1);
    containers.add(mockContainer2);
    containers.add(mockContainer3);
    containers.add(mockContainer4);
    scheduler.onContainersAllocated(containers);
    drainableAppCallback.drain();
    // first container allocated
    verify(mockApp).taskAllocated(mockTask1, mockCookie1, mockContainer1);
    verify(mockApp).taskAllocated(mockTask2, mockCookie2, mockContainer2);
    verify(mockApp).taskAllocated(mockTask3, mockCookie3, mockContainer3);
    // no other allocations returned
    verify(mockApp, times(3)).taskAllocated(any(), any(), (Container) any());
    verify(mockRMClient).removeContainerRequest(request1);
    verify(mockRMClient).removeContainerRequest(request2);
    verify(mockRMClient).removeContainerRequest(request3);
    // verify unwanted container released
    verify(mockRMClient).releaseAssignedContainer(mockCId4);

    // deallocate allocated task
    assertTrue(scheduler.deallocateTask(mockTask1, true, null, null));
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(mockCId1);
    verify(mockRMClient).releaseAssignedContainer(mockCId1);
    // deallocate allocated container
    assertEquals(mockTask2, scheduler.deallocateContainer(mockCId2));
    drainableAppCallback.drain();
    verify(mockRMClient).releaseAssignedContainer(mockCId2);
    verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any());

    List<ContainerStatus> statuses = new ArrayList<>();
    ContainerStatus mockStatus1 = mock(ContainerStatus.class);
    when(mockStatus1.getContainerId()).thenReturn(mockCId1);
    statuses.add(mockStatus1);
    ContainerStatus mockStatus2 = mock(ContainerStatus.class);
    when(mockStatus2.getContainerId()).thenReturn(mockCId2);
    statuses.add(mockStatus2);
    ContainerStatus mockStatus3 = mock(ContainerStatus.class);
    when(mockStatus3.getContainerId()).thenReturn(mockCId3);
    statuses.add(mockStatus3);
    ContainerStatus mockStatus4 = mock(ContainerStatus.class);
    when(mockStatus4.getContainerId()).thenReturn(mockCId4);
    statuses.add(mockStatus4);

    scheduler.onContainersCompleted(statuses);
    drainableAppCallback.drain();
    // released container status returned
    verify(mockApp).containerCompleted(mockTask1, mockStatus1);
    verify(mockApp).containerCompleted(mockTask2, mockStatus2);
    // currently allocated container status returned and not released
    verify(mockApp).containerCompleted(mockTask3, mockStatus3);
    // no other statuses returned
    verify(mockApp, times(3)).containerCompleted(any(), (ContainerStatus) any());
    verify(mockRMClient, times(3)).releaseAssignedContainer((ContainerId) any());

    // verify blacklisting
    verify(mockRMClient, times(0)).updateBlacklist(anyListOf(String.class), anyListOf(String.class));
    String badHost = "host6";
    NodeId badNodeId = NodeId.newInstance(badHost, 1);
    scheduler.blacklistNode(badNodeId);
    List<String> badNodeList = Collections.singletonList(badHost);
    verify(mockRMClient, times(1)).updateBlacklist(eq(badNodeList), isNull(List.class));
    Object mockTask4 = new MockTask("task4");
    Object mockCookie4 = new Object();
    scheduler.allocateTask(mockTask4, mockCapability, null, null, mockPriority, null, mockCookie4);
    drainableAppCallback.drain();
    verify(mockRMClient, times(5)).addContainerRequest(requestCaptor.capture());
    ContainerId mockCId5 = ContainerId.newContainerId(attemptId, 5);
    Container mockContainer5 = Container.newInstance(mockCId5, badNodeId, null, mockCapability, mockPriority,
            null);
    containers.clear();
    containers.add(mockContainer5);
    scheduler.onContainersAllocated(containers);
    drainableAppCallback.drain();
    // no new allocation
    verify(mockApp, times(3)).taskAllocated(any(), any(), (Container) any());
    // verify blacklisted container released
    verify(mockRMClient).releaseAssignedContainer(mockCId5);
    verify(mockRMClient, times(4)).releaseAssignedContainer((ContainerId) any());
    // verify request added back
    verify(mockRMClient, times(6)).addContainerRequest(requestCaptor.capture());
    NodeId host6 = NodeId.newInstance("host6", 6);
    ContainerId mockCId6 = ContainerId.newContainerId(attemptId, 6);
    Container mockContainer6 = Container.newInstance(mockCId6, host6, null, mockCapability, mockPriority, null);
    containers.clear();
    containers.add(mockContainer6);
    scheduler.onContainersAllocated(containers);
    drainableAppCallback.drain();
    // new allocation
    verify(mockApp, times(4)).taskAllocated(any(), any(), (Container) any());
    verify(mockApp).taskAllocated(mockTask4, mockCookie4, mockContainer6);
    // deallocate allocated task
    assertTrue(scheduler.deallocateTask(mockTask4, true, null, null));
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(mockCId6);
    verify(mockRMClient).releaseAssignedContainer(mockCId6);
    verify(mockRMClient, times(5)).releaseAssignedContainer((ContainerId) any());
    // test unblacklist
    scheduler.unblacklistNode(badNodeId);
    verify(mockRMClient, times(1)).updateBlacklist(isNull(List.class), eq(badNodeList));
    assertEquals(0, scheduler.getNumBlacklistedNodes());

    float progress = 0.5f;
    when(mockApp.getProgress()).thenReturn(progress);
    assertEquals(progress, scheduler.getProgress(), 0);

    // check duplicate allocation request
    scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1);
    drainableAppCallback.drain();
    verify(mockRMClient, times(7)).addContainerRequest(any(TaskRequest.class));
    verify(mockRMClient, times(6)).removeContainerRequest(any(TaskRequest.class));
    scheduler.allocateTask(mockTask1, mockCapability, hosts, racks, mockPriority, null, mockCookie1);
    drainableAppCallback.drain();
    // old request removed and new one added
    verify(mockRMClient, times(7)).removeContainerRequest(any(TaskRequest.class));
    verify(mockRMClient, times(8)).addContainerRequest(any(TaskRequest.class));
    assertFalse(scheduler.deallocateTask(mockTask1, true, null, null));

    List<NodeReport> mockUpdatedNodes = mock(List.class);
    scheduler.onNodesUpdated(mockUpdatedNodes);
    drainableAppCallback.drain();
    verify(mockApp).nodesUpdated(mockUpdatedNodes);

    ArgumentCaptor<String> argumentCaptor = ArgumentCaptor.forClass(String.class);
    Exception mockException = new IOException("mockexception");
    scheduler.onError(mockException);
    drainableAppCallback.drain();
    verify(mockApp).reportError(eq(YarnTaskSchedulerServiceError.RESOURCEMANAGER_ERROR),
            argumentCaptor.capture(), any(DagInfo.class));
    assertTrue(argumentCaptor.getValue().contains("mockexception"));

    scheduler.onShutdownRequest();
    drainableAppCallback.drain();
    verify(mockApp).appShutdownRequested();

    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    verify(mockRMClient).stop();
}

From source file:org.apache.tez.dag.app.rm.TestDagAwareYarnTaskScheduler.java

License:Apache License

@Test(timeout = 30000)
public void testSimpleReuseLocalMatching() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, false);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock);

    scheduler.initialize();/*w w  w.  jav a 2  s  .  c  om*/
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
            regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
            regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, "host3");
    MockTaskInfo taskv0t2 = new MockTaskInfo("taskv0t2", priorityv0, hostsv0t0);
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, hostsv0t0);
    MockTaskInfo taskv1t1 = new MockTaskInfo("taskv1t1", priorityv1, hostsv0t0);

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    taskRequestCaptor.scheduleTask(taskv0t1);
    TaskRequest reqv0t2 = taskRequestCaptor.scheduleTask(taskv0t2);
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0);
    taskRequestCaptor.scheduleTask(taskv1t1);

    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv0t0.capability, priorityv0, null);

    // allocate one container at v0 priority
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);

    // finish v0t0 successfully, verify v0t1 is skipped and v0t2 instead is assigned to the container
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv0t2.task, taskv0t2.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t2);

    // finish v0t2 successfully, verify v1t0 is assigned to the same container
    assertTrue(scheduler.deallocateTask(taskv0t2.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv1t0);

    // fail v1t0 and verify container is released instead of reused for v1t1
    assertTrue(scheduler.deallocateTask(taskv1t0.task, false, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(cid1);

    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    verify(mockRMClient).stop();
}

From source file:org.apache.tez.dag.app.rm.TestDagAwareYarnTaskScheduler.java

License:Apache License

@Test(timeout = 30000)
public void testSimpleReuseRackMatching() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock);

    scheduler.initialize();//from w  w w . ja va2 s  . co m
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
            regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
            regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, "host2");
    MockTaskInfo taskv0t2 = new MockTaskInfo("taskv0t2", priorityv0, "host4", "/somerack");
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, "host1");
    MockTaskInfo taskv1t1 = new MockTaskInfo("taskv1t1", priorityv1, "host5");

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    TaskRequest reqv0t1 = taskRequestCaptor.scheduleTask(taskv0t1);
    taskRequestCaptor.scheduleTask(taskv0t2);
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0);
    taskRequestCaptor.scheduleTask(taskv1t1);

    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv0t0.capability, priorityv0, null);

    // allocate one container at v0 priority
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);

    // finish v0t0 successfully, verify v0t1 is skipped and v1t0 assigned instead
    // since host locality is preferred to rack locality and lower priority vertex
    // is not blocked by higher priority vertex
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv1t0);

    // finish v1t0 successfully, verify v0t1 is assigned
    assertTrue(scheduler.deallocateTask(taskv1t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv0t1.task, taskv0t1.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t1);

    // fail v0t1 and verify container is released instead of reused for v1t1
    assertTrue(scheduler.deallocateTask(taskv0t1.task, false, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(cid1);

    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    verify(mockRMClient).stop();
}

From source file:org.apache.tez.dag.app.rm.TestDagAwareYarnTaskScheduler.java

License:Apache License

@Test(timeout = 30000)
public void testSimpleReuseAnyMatching() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock);

    scheduler.initialize();/*from ww w.j  av  a2  s .c o m*/
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
            regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
            regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, "host2");
    MockTaskInfo taskv0t2 = new MockTaskInfo("taskv0t2", priorityv0, "host4", "/rack4");
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, "host1");
    MockTaskInfo taskv1t1 = new MockTaskInfo("taskv1t1", priorityv1, "host6", "/rack6");

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    TaskRequest reqv0t1 = taskRequestCaptor.scheduleTask(taskv0t1);
    TaskRequest reqv0t2 = taskRequestCaptor.scheduleTask(taskv0t2);
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0);
    taskRequestCaptor.scheduleTask(taskv1t1);

    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv0t0.capability, priorityv0, null);

    // allocate one container at v0 priority
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);

    // finish v0t0 successfully, verify v0t1 is skipped and v1t0 assigned instead
    // since host locality is preferred to rack locality and lower priority vertex
    // is not blocked by higher priority vertex
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv1t0);

    // finish v1t0 successfully, verify v0t1 is assigned
    assertTrue(scheduler.deallocateTask(taskv1t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv0t1.task, taskv0t1.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t1);

    // finish v0t1 successfully, verify v0t2 is assigned
    assertTrue(scheduler.deallocateTask(taskv0t1.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv0t2.task, taskv0t2.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t2);

    // fail v0t2 and verify container is released instead of reused for v1t1
    assertTrue(scheduler.deallocateTask(taskv0t2.task, false, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(cid1);

    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    verify(mockRMClient).stop();
}

From source file:org.apache.tez.dag.app.rm.TestDagAwareYarnTaskScheduler.java

License:Apache License

@Test(timeout = 30000)
public void testReuseWithAffinity() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock);

    scheduler.initialize();//from  w w  w .j  a  v a  2 s . c  o  m
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
            regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
            regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, hostsv0t0);

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    taskRequestCaptor.scheduleTask(taskv0t1);

    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv0t0.capability, priorityv0, null);

    // allocate one container at v0 priority
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);

    // add a new request for this container
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, "host1");
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0, cid1);

    // finish v0t0 successfully, verify v0t1 is skipped even though it is node-local
    // and v1t0 assigned instead for affinity
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv1t0);

    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    verify(mockRMClient).stop();
}

From source file:org.apache.tez.dag.app.rm.TestDagAwareYarnTaskScheduler.java

License:Apache License

@Test(timeout = 30000)
public void testReuseVertexDescendants() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);

    // vertex 0 and vertex 2 are root vertices and vertex 1 is a child of vertex 0
    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(3);
    when(mockDagInfo.getVertexDescendants(0)).thenReturn(BitSet.valueOf(new long[] { 0x2 }));
    when(mockDagInfo.getVertexDescendants(1)).thenReturn(new BitSet());
    when(mockDagInfo.getVertexDescendants(2)).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock);

    scheduler.initialize();//  w ww  . j a  va 2  s .  c  om
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
            regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
            regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    Priority priorityv0 = Priority.newInstance(1);
    Priority priorityv1 = Priority.newInstance(2);
    Priority priorityv2 = Priority.newInstance(3);
    String[] hostsv0t0 = { "host1", "host2" };
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv0t0.task)).thenReturn(0);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, "host3");
    when(mockApp.getVertexIndexForTask(taskv0t1.task)).thenReturn(0);
    MockTaskInfo taskv1t0 = new MockTaskInfo("taskv1t0", priorityv1, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv1t0.task)).thenReturn(1);
    MockTaskInfo taskv2t0 = new MockTaskInfo("taskv2t0", priorityv2, hostsv0t0);
    when(mockApp.getVertexIndexForTask(taskv2t0.task)).thenReturn(2);
    MockTaskInfo taskv2t1 = new MockTaskInfo("taskv2t1", priorityv2, "host3");
    when(mockApp.getVertexIndexForTask(taskv2t1.task)).thenReturn(2);

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    TaskRequest reqv0t1 = taskRequestCaptor.scheduleTask(taskv0t1);
    TaskRequest reqv1t0 = taskRequestCaptor.scheduleTask(taskv1t0);
    TaskRequest reqv2t0 = taskRequestCaptor.scheduleTask(taskv2t0);
    taskRequestCaptor.scheduleTask(taskv2t1);

    NodeId host1 = NodeId.newInstance("host1", 1);
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    Container container1 = Container.newInstance(cid1, host1, null, taskv0t0.capability, priorityv0, null);

    // allocate one container at v0 priority
    scheduler.onContainersAllocated(Collections.singletonList(container1));
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);

    // finish v0t0 successfully, verify v1t0 is skipped and v2t0 assigned instead
    // since host locality beats rack locality for unblocked vertex v2 and
    // v1 is blocked by pending v0 request
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv2t0.task, taskv2t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv2t0);

    // finish v2t0 successfully, verify v0t1 is assigned since it is higher
    // priority than v2
    assertTrue(scheduler.deallocateTask(taskv2t0.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv0t1.task, taskv0t1.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t1);

    // finish v2t0 successfully, verify v1t0 is assigned since it is now unblocked
    assertTrue(scheduler.deallocateTask(taskv0t1.task, true, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));
    verify(mockApp).taskAllocated(taskv1t0.task, taskv1t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv1t0);

    // fail v1t0 and verify container is released instead of reused for v2t1
    assertTrue(scheduler.deallocateTask(taskv1t0.task, false, null, null));
    clock.incrementTime(10000);
    drainableAppCallback.drain();
    verify(mockApp).containerBeingReleased(cid1);
    verify(mockRMClient).releaseAssignedContainer(cid1);

    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    verify(mockRMClient).stop();
}

From source file:org.apache.tez.dag.app.rm.TestDagAwareYarnTaskScheduler.java

License:Apache License

@Test(timeout = 30000)
public void testSessionContainers() throws Exception {
    AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest());

    String appHost = "host";
    int appPort = 0;
    String appUrl = "url";

    Configuration conf = new Configuration();
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, true);
    conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false);
    conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 4000);
    conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 5000);
    conf.setInt(TezConfiguration.TEZ_AM_SESSION_MIN_HELD_CONTAINERS, 5);

    DagInfo mockDagInfo = mock(DagInfo.class);
    when(mockDagInfo.getTotalVertices()).thenReturn(10);
    when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet());
    TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf);
    when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo);
    when(mockApp.isSession()).thenReturn(true);
    TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp);

    MockClock clock = new MockClock(1000);
    NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock);

    scheduler.initialize();/*from   w ww. ja  v a  2  s . co  m*/
    drainableAppCallback.drain();

    scheduler.start();
    drainableAppCallback.drain();
    verify(mockRMClient).start();
    verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl);
    RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse();
    verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(),
            regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(),
            regResponse.getQueue());

    assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount());

    final String rack1 = "/r1";
    final String rack2 = "/r2";
    final String rack3 = "/r3";
    final String node1Rack1 = "n1r1";
    final String node2Rack1 = "n2r1";
    final String node1Rack2 = "n1r2";
    final String node2Rack2 = "n2r2";
    final String node1Rack3 = "n1r3";
    MockDNSToSwitchMapping.addRackMapping(node1Rack1, rack1);
    MockDNSToSwitchMapping.addRackMapping(node2Rack1, rack1);
    MockDNSToSwitchMapping.addRackMapping(node1Rack2, rack2);
    MockDNSToSwitchMapping.addRackMapping(node2Rack2, rack2);
    MockDNSToSwitchMapping.addRackMapping(node1Rack3, rack3);

    Priority priorityv0 = Priority.newInstance(1);
    MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, node1Rack1, rack1);
    MockTaskInfo taskv0t1 = new MockTaskInfo("taskv0t1", priorityv0, node2Rack1, rack1);
    MockTaskInfo taskv0t2 = new MockTaskInfo("taskv0t2", priorityv0, node1Rack1, rack1);
    MockTaskInfo taskv0t3 = new MockTaskInfo("taskv0t3", priorityv0, node2Rack1, rack1);
    MockTaskInfo taskv0t4 = new MockTaskInfo("taskv0t4", priorityv0, node1Rack2, rack2);
    MockTaskInfo taskv0t5 = new MockTaskInfo("taskv0t5", priorityv0, node2Rack2, rack2);
    MockTaskInfo taskv0t6 = new MockTaskInfo("taskv0t6", priorityv0, node1Rack3, rack3);

    TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback);
    TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0);
    TaskRequest reqv0t1 = taskRequestCaptor.scheduleTask(taskv0t1);
    TaskRequest reqv0t2 = taskRequestCaptor.scheduleTask(taskv0t2);
    TaskRequest reqv0t3 = taskRequestCaptor.scheduleTask(taskv0t3);
    TaskRequest reqv0t4 = taskRequestCaptor.scheduleTask(taskv0t4);
    TaskRequest reqv0t5 = taskRequestCaptor.scheduleTask(taskv0t5);
    TaskRequest reqv0t6 = taskRequestCaptor.scheduleTask(taskv0t6);

    List<Container> containers = new ArrayList<>();
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1);
    ContainerId cid1 = ContainerId.newContainerId(attemptId, 1);
    NodeId n1r1 = NodeId.newInstance(node1Rack1, 1);
    Container container1 = Container.newInstance(cid1, n1r1, null, taskv0t0.capability, priorityv0, null);
    containers.add(container1);
    ContainerId cid2 = ContainerId.newContainerId(attemptId, 2);
    NodeId n2r1 = NodeId.newInstance(node2Rack1, 1);
    Container container2 = Container.newInstance(cid2, n2r1, null, taskv0t1.capability, priorityv0, null);
    containers.add(container2);
    ContainerId cid3 = ContainerId.newContainerId(attemptId, 3);
    Container container3 = Container.newInstance(cid3, n1r1, null, taskv0t2.capability, priorityv0, null);
    containers.add(container3);
    ContainerId cid4 = ContainerId.newContainerId(attemptId, 4);
    Container container4 = Container.newInstance(cid4, n2r1, null, taskv0t3.capability, priorityv0, null);
    containers.add(container4);
    ContainerId cid5 = ContainerId.newContainerId(attemptId, 5);
    NodeId n1r2 = NodeId.newInstance(node1Rack2, 1);
    Container container5 = Container.newInstance(cid5, n1r2, null, taskv0t4.capability, priorityv0, null);
    containers.add(container5);
    ContainerId cid6 = ContainerId.newContainerId(attemptId, 6);
    NodeId n2r2 = NodeId.newInstance(node2Rack2, 1);
    Container container6 = Container.newInstance(cid6, n2r2, null, taskv0t5.capability, priorityv0, null);
    containers.add(container6);
    ContainerId cid7 = ContainerId.newContainerId(attemptId, 7);
    NodeId n1r3 = NodeId.newInstance(node1Rack3, 1);
    Container container7 = Container.newInstance(cid7, n1r3, null, taskv0t6.capability, priorityv0, null);
    containers.add(container7);

    scheduler.onContainersAllocated(containers);
    drainableAppCallback.drain();
    verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1);
    verify(mockRMClient).removeContainerRequest(reqv0t0);
    verify(mockApp).taskAllocated(taskv0t1.task, taskv0t1.cookie, container2);
    verify(mockRMClient).removeContainerRequest(reqv0t1);
    verify(mockApp).taskAllocated(taskv0t2.task, taskv0t2.cookie, container3);
    verify(mockRMClient).removeContainerRequest(reqv0t2);
    verify(mockApp).taskAllocated(taskv0t3.task, taskv0t3.cookie, container4);
    verify(mockRMClient).removeContainerRequest(reqv0t3);
    verify(mockApp).taskAllocated(taskv0t4.task, taskv0t4.cookie, container5);
    verify(mockRMClient).removeContainerRequest(reqv0t4);
    verify(mockApp).taskAllocated(taskv0t5.task, taskv0t5.cookie, container6);
    verify(mockRMClient).removeContainerRequest(reqv0t5);
    verify(mockApp).taskAllocated(taskv0t6.task, taskv0t6.cookie, container7);
    verify(mockRMClient).removeContainerRequest(reqv0t6);

    clock.incrementTime(10000);
    drainableAppCallback.drain();
    assertTrue(scheduler.deallocateTask(taskv0t0.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t1.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t2.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t3.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t4.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t5.task, true, null, null));
    assertTrue(scheduler.deallocateTask(taskv0t6.task, true, null, null));
    verify(mockApp, never()).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, never()).releaseAssignedContainer(any(ContainerId.class));

    // verify only two of the containers were released after idle expiration
    // and the rest were spread across the nodes and racks
    clock.incrementTime(5000);
    drainableAppCallback.drain();
    verify(mockApp, times(2)).containerBeingReleased(any(ContainerId.class));
    verify(mockRMClient, times(2)).releaseAssignedContainer(any(ContainerId.class));
    Set<String> hosts = new HashSet<>();
    Set<String> racks = new HashSet<>();
    for (HeldContainer hc : scheduler.getSessionContainers()) {
        hosts.add(hc.getHost());
        racks.add(hc.getRack());
    }
    assertEquals(5, hosts.size());
    assertEquals(3, racks.size());
    assertTrue(hosts.contains(node1Rack1));
    assertTrue(hosts.contains(node2Rack1));
    assertTrue(hosts.contains(node1Rack2));
    assertTrue(hosts.contains(node2Rack2));
    assertTrue(hosts.contains(node1Rack3));
    assertTrue(racks.contains(rack1));
    assertTrue(racks.contains(rack2));
    assertTrue(racks.contains(rack3));

    String appMsg = "success";
    AppFinalStatus finalStatus = new AppFinalStatus(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    when(mockApp.getFinalAppStatus()).thenReturn(finalStatus);
    scheduler.shutdown();
    drainableAppCallback.drain();
    verify(mockRMClient).unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, appMsg, appUrl);
    verify(mockRMClient).stop();
}