Example usage for org.apache.hadoop.yarn.api.records Resource newInstance

List of usage examples for org.apache.hadoop.yarn.api.records Resource newInstance

Introduction

In this page you can find the example usage for org.apache.hadoop.yarn.api.records Resource newInstance.

Prototype

@Public
    @Stable
    public static Resource newInstance(long memory, int vCores) 

Source Link

Usage

From source file:org.apache.myriad.TestObjectFactory.java

License:Apache License

public static RMNode getRMNode(String host, int port, int memory, int cores) {
    Resource resource = Resource.newInstance(memory, cores);
    return getRMNode(host, port, resource);
}

From source file:org.apache.reef.runtime.yarn.client.YarnJobSubmissionHandler.java

License:Apache License

@Override
public void onNext(final ClientRuntimeProtocol.JobSubmissionProto jobSubmissionProto) {

    LOG.log(Level.FINEST, "Submitting job with ID [{0}]", jobSubmissionProto.getIdentifier());

    try {/*w w  w.j a  v a  2 s .com*/

        LOG.log(Level.FINE, "Requesting Application ID from YARN.");

        final YarnClientApplication yarnClientApplication = this.yarnClient.createApplication();
        final GetNewApplicationResponse applicationResponse = yarnClientApplication.getNewApplicationResponse();

        final ApplicationSubmissionContext applicationSubmissionContext = yarnClientApplication
                .getApplicationSubmissionContext();

        final ApplicationId applicationId = applicationSubmissionContext.getApplicationId();

        LOG.log(Level.FINEST, "YARN Application ID: {0}", applicationId);

        // set the application name
        applicationSubmissionContext.setApplicationName("reef-job-" + jobSubmissionProto.getIdentifier());

        LOG.log(Level.FINE, "Assembling submission JAR for the Driver.");

        final Path submissionFolder = new Path(
                "/tmp/" + this.filenames.getJobFolderPrefix() + applicationId.getId() + "/");

        final Configuration driverConfiguration = makeDriverConfiguration(jobSubmissionProto, submissionFolder);

        final File jobSubmissionFile = this.jobJarMaker.createJobSubmissionJAR(jobSubmissionProto,
                driverConfiguration);

        final Path uploadedJobJarPath = this.uploadToJobFolder(jobSubmissionFile, submissionFolder);

        final Map<String, LocalResource> resources = new HashMap<>(1);
        resources.put(this.filenames.getREEFFolderName(), this.makeLocalResourceForJarFile(uploadedJobJarPath));

        // SET MEMORY RESOURCE
        final int amMemory = getMemory(jobSubmissionProto,
                applicationResponse.getMaximumResourceCapability().getMemory());
        applicationSubmissionContext.setResource(Resource.newInstance(amMemory, 1));

        // SET EXEC COMMAND
        final List<String> launchCommand = new JavaLaunchCommandBuilder()
                .setErrorHandlerRID(jobSubmissionProto.getRemoteId())
                .setLaunchID(jobSubmissionProto.getIdentifier())
                .setConfigurationFileName(this.filenames.getDriverConfigurationPath())
                .setClassPath(this.classpath.getDriverClasspath()).setMemory(amMemory)
                .setStandardOut(ApplicationConstants.LOG_DIR_EXPANSION_VAR + "/"
                        + this.filenames.getDriverStdoutFileName())
                .setStandardErr(ApplicationConstants.LOG_DIR_EXPANSION_VAR + "/"
                        + this.filenames.getDriverStderrFileName())
                .build();

        applicationSubmissionContext
                .setAMContainerSpec(YarnTypes.getContainerLaunchContext(launchCommand, resources));

        applicationSubmissionContext.setPriority(getPriority(jobSubmissionProto));

        // Set the queue to which this application is to be submitted in the RM
        applicationSubmissionContext.setQueue(getQueue(jobSubmissionProto, "default"));
        LOG.log(Level.INFO, "Submitting REEF Application to YARN. ID: {0}", applicationId);

        if (LOG.isLoggable(Level.FINEST)) {
            LOG.log(Level.FINEST, "REEF app command: {0}", StringUtils.join(launchCommand, ' '));
        }

        // TODO: this is currently being developed on a hacked 2.4.0 bits, should be 2.4.1
        final String minVersionKeepContainerOptionAvailable = "2.4.0";

        // when supported, set KeepContainersAcrossApplicationAttempts to be true
        // so that when driver (AM) crashes, evaluators will still be running and we can recover later.
        if (YarnTypes.isAtOrAfterVersion(minVersionKeepContainerOptionAvailable)) {
            LOG.log(Level.FINE,
                    "Hadoop version is {0} or after with KeepContainersAcrossApplicationAttempts supported, will set it to true.",
                    minVersionKeepContainerOptionAvailable);

            applicationSubmissionContext.setKeepContainersAcrossApplicationAttempts(true);
        }

        this.yarnClient.submitApplication(applicationSubmissionContext);

    } catch (final YarnException | IOException e) {
        throw new RuntimeException("Unable to submit Driver to YARN.", e);
    }
}

From source file:org.apache.reef.runtime.yarn.client.YarnSubmissionHelper.java

License:Apache License

/**
 * Set the amount of memory to be allocated to the Driver.
 * @param megabytes//from ww  w  . jav  a 2 s .  c o  m
 * @return
 */
public YarnSubmissionHelper setDriverMemory(final int megabytes) {
    applicationSubmissionContext.setResource(Resource.newInstance(getMemory(megabytes), 1));
    return this;
}

From source file:org.apache.samza.job.yarn.refactor.YarnClusterResourceManager.java

License:Apache License

/**
 * Request resources for running container processes.
 *///  w  w  w .j  a  v  a2s .  com
@Override
public void requestResources(SamzaResourceRequest resourceRequest) {
    final int DEFAULT_PRIORITY = 0;
    log.info("Requesting resources on  " + resourceRequest.getPreferredHost() + " for container "
            + resourceRequest.getContainerID());

    int memoryMb = resourceRequest.getMemoryMB();
    int cpuCores = resourceRequest.getNumCores();
    String preferredHost = resourceRequest.getPreferredHost();
    Resource capability = Resource.newInstance(memoryMb, cpuCores);
    Priority priority = Priority.newInstance(DEFAULT_PRIORITY);

    AMRMClient.ContainerRequest issuedRequest;

    if (preferredHost.equals("ANY_HOST")) {
        log.info("Making a request for ANY_HOST " + preferredHost);
        issuedRequest = new AMRMClient.ContainerRequest(capability, null, null, priority);
    } else {
        log.info("Making a preferred host request on " + preferredHost);
        issuedRequest = new AMRMClient.ContainerRequest(capability, new String[] { preferredHost }, null,
                priority);
    }
    //ensure that updating the state and making the request are done atomically.
    synchronized (lock) {
        requestsMap.put(resourceRequest, issuedRequest);
        amClient.addContainerRequest(issuedRequest);
    }
}

From source file:org.apache.samza.job.yarn.SamzaContainerRequest.java

License:Apache License

public SamzaContainerRequest(int memoryMb, int cpuCores, int priority, int expectedContainerId,
        String preferredHost) {//  www.j  av a 2 s . co m
    this.capability = Resource.newInstance(memoryMb, cpuCores);
    this.priority = Priority.newInstance(priority);
    this.expectedContainerId = expectedContainerId;
    if (preferredHost == null) {
        this.preferredHost = ANY_HOST;
        this.issuedRequest = new AMRMClient.ContainerRequest(capability, null, null, this.priority);
    } else {
        this.preferredHost = preferredHost;
        this.issuedRequest = new AMRMClient.ContainerRequest(capability, new String[] { this.preferredHost },
                null, this.priority);
    }

    this.requestTimestamp = System.currentTimeMillis();
}

From source file:org.apache.samza.job.yarn.TestYarnClusterResourceManager.java

License:Apache License

@Test
public void testErrorInStartContainerShouldUpdateState() {
    // create mocks
    final int samzaContainerId = 1;
    YarnConfiguration yarnConfiguration = mock(YarnConfiguration.class);
    SamzaAppMasterMetrics metrics = mock(SamzaAppMasterMetrics.class);
    Config config = mock(Config.class);
    AMRMClientAsync asyncClient = mock(AMRMClientAsync.class);
    YarnAppState yarnAppState = new YarnAppState(0, mock(ContainerId.class), "host", 8080, 8081);
    SamzaYarnAppMasterLifecycle lifecycle = mock(SamzaYarnAppMasterLifecycle.class);
    SamzaYarnAppMasterService service = mock(SamzaYarnAppMasterService.class);
    NMClientAsync asyncNMClient = mock(NMClientAsync.class);
    ClusterResourceManager.Callback callback = mock(ClusterResourceManager.Callback.class);

    // start the cluster manager
    YarnClusterResourceManager yarnClusterResourceManager = new YarnClusterResourceManager(asyncClient,
            asyncNMClient, callback, yarnAppState, lifecycle, service, metrics, yarnConfiguration, config);

    yarnAppState.pendingProcessors.put(String.valueOf(samzaContainerId),
            new YarnContainer(Container.newInstance(
                    ContainerId.newContainerId(
                            ApplicationAttemptId.newInstance(ApplicationId.newInstance(10000l, 1), 1), 1),
                    NodeId.newInstance("host1", 8088), "http://host1", Resource.newInstance(1024, 1),
                    Priority.newInstance(1),
                    Token.newInstance("id".getBytes(), "read", "password".getBytes(), "service"))));

    yarnClusterResourceManager.start();//from  w w w.j av a  2s  . com
    assertEquals(1, yarnAppState.pendingProcessors.size());

    yarnClusterResourceManager.onStartContainerError(
            ContainerId.newContainerId(
                    ApplicationAttemptId.newInstance(ApplicationId.newInstance(10000l, 1), 1), 1),
            new Exception());

    assertEquals(0, yarnAppState.pendingProcessors.size());
    verify(callback, times(1)).onStreamProcessorLaunchFailure(anyObject(), any(Exception.class));
}

From source file:org.apache.samza.job.yarn.YarnClusterResourceManager.java

License:Apache License

/**
 * Request resources for running container processes.
 *///from w ww .  j a v  a  2s .c  om
@Override
public void requestResources(SamzaResourceRequest resourceRequest) {
    String processorId = resourceRequest.getProcessorId();
    String requestId = resourceRequest.getRequestId();
    String preferredHost = resourceRequest.getPreferredHost();
    int memoryMb = resourceRequest.getMemoryMB();
    int cpuCores = resourceRequest.getNumCores();
    Resource capability = Resource.newInstance(memoryMb, cpuCores);
    String nodeLabelsExpression = yarnConfig.getContainerLabel();

    AMRMClient.ContainerRequest issuedRequest;

    /*
     * Yarn enforces these two checks:
     *   1. ANY_HOST requests should always be made with relax-locality = true
     *   2. A request with relax-locality = false should not be in the same priority as another with relax-locality = true
     *
     * Since the Samza AM makes preferred-host requests with relax-locality = false, it follows that ANY_HOST requests
     * should specify a different priority-level. We can safely set priority of preferred-host requests to be higher than
     * any-host requests since data-locality is critical.
     */
    if (preferredHost.equals("ANY_HOST")) {
        Priority priority = Priority.newInstance(ANY_HOST_PRIORITY);
        boolean relaxLocality = true;
        log.info(
                "Requesting resources for Processor ID: {} on nodes: {} on racks: {} with capability: {}, priority: {}, relaxLocality: {}, nodeLabelsExpression: {}",
                processorId, null, null, capability, priority, relaxLocality, nodeLabelsExpression);
        issuedRequest = new AMRMClient.ContainerRequest(capability, null, null, priority, relaxLocality,
                nodeLabelsExpression);
    } else {
        String[] nodes = { preferredHost };
        Priority priority = Priority.newInstance(PREFERRED_HOST_PRIORITY);
        boolean relaxLocality = false;
        log.info(
                "Requesting resources for Processor ID: {} on nodes: {} on racks: {} with capability: {}, priority: {}, relaxLocality: {}, nodeLabelsExpression: {}",
                processorId, Arrays.toString(nodes), null, capability, priority, relaxLocality,
                nodeLabelsExpression);
        issuedRequest = new AMRMClient.ContainerRequest(capability, nodes, null, priority, relaxLocality,
                nodeLabelsExpression);
    }
    // ensure that updating the state and making the request are done atomically.
    synchronized (lock) {
        requestsMap.put(resourceRequest, issuedRequest);
        amClient.addContainerRequest(issuedRequest);
    }
}

From source file:org.apache.sysml.yarn.ropt.YarnClusterAnalyzer.java

License:Apache License

/**
 * Analyzes properties of Yarn cluster and Hadoop configurations.
 * /*from w w  w  .  j  a  v  a2  s. com*/
 * @param yarnClient hadoop yarn client
 * @param conf hadoop yarn configuration
 * @param verbose output info to standard output
 */
public static void analyzeYarnCluster(YarnClient yarnClient, YarnConfiguration conf, boolean verbose) {
    try {
        List<NodeReport> nodesReport = yarnClient.getNodeReports();
        if (verbose)
            System.out.println("There are " + nodesReport.size() + " nodes in the cluster");
        if (nodesReport.isEmpty())
            throw new YarnException("There are zero available nodes in the yarn cluster");

        nodesMaxPhySorted = new ArrayList<>(nodesReport.size());
        clusterTotalMem = 0;
        clusterTotalCores = 0;
        clusterTotalNodes = 0;
        minimumMRContainerPhyMB = -1;
        for (NodeReport node : nodesReport) {
            Resource resource = node.getCapability();
            Resource used = node.getUsed();
            if (used == null)
                used = Resource.newInstance(0, 0);
            int mb = resource.getMemory();
            int cores = resource.getVirtualCores();
            if (mb <= 0)
                throw new YarnException("A node has non-positive memory " + mb);

            int myMinMRPhyMB = mb / cores / CPU_HYPER_FACTOR;
            if (minimumMRContainerPhyMB < myMinMRPhyMB)
                minimumMRContainerPhyMB = myMinMRPhyMB; // minimumMRContainerPhyMB needs to be the largest among the mins

            clusterTotalMem += (long) mb * 1024 * 1024;
            nodesMaxPhySorted.add((long) mb * 1024 * 1024);
            clusterTotalCores += cores;
            clusterTotalNodes++;
            if (verbose)
                System.out.println("\t" + node.getNodeId() + " has " + mb + " MB (" + used.getMemory()
                        + " MB used) memory and " + resource.getVirtualCores() + " (" + used.getVirtualCores()
                        + " used) cores");

        }
        Collections.sort(nodesMaxPhySorted, Collections.reverseOrder());

        nodesMaxBudgetSorted = new ArrayList<>(nodesMaxPhySorted.size());
        for (int i = 0; i < nodesMaxPhySorted.size(); i++)
            nodesMaxBudgetSorted.add(ResourceOptimizer.phyToBudget(nodesMaxPhySorted.get(i)));

        _remotePar = nodesReport.size();
        if (_remotePar == 0)
            throw new YarnException("There are no available nodes in the yarn cluster");

        // Now get the default cluster settings
        _remoteMRSortMem = (1024 * 1024) * conf.getLong(MRConfigurationNames.MR_TASK_IO_SORT_MB, 100); //100MB

        //handle jvm max mem (map mem budget is relevant for map-side distcache and parfor)
        //(for robustness we probe both: child and map configuration parameters)
        String javaOpts1 = conf.get(MRConfigurationNames.MR_CHILD_JAVA_OPTS); //internally mapred/mapreduce synonym
        String javaOpts2 = conf.get(MRConfigurationNames.MR_MAP_JAVA_OPTS, null); //internally mapred/mapreduce synonym
        String javaOpts3 = conf.get(MRConfigurationNames.MR_REDUCE_JAVA_OPTS, null); //internally mapred/mapreduce synonym
        if (javaOpts2 != null) //specific value overrides generic
            _remoteJVMMaxMemMap = extractMaxMemoryOpt(javaOpts2);
        else
            _remoteJVMMaxMemMap = extractMaxMemoryOpt(javaOpts1);
        if (javaOpts3 != null) //specific value overrides generic
            _remoteJVMMaxMemReduce = extractMaxMemoryOpt(javaOpts3);
        else
            _remoteJVMMaxMemReduce = extractMaxMemoryOpt(javaOpts1);

        //HDFS blocksize
        String blocksize = conf.get(MRConfigurationNames.DFS_BLOCKSIZE, "134217728");
        _blocksize = Long.parseLong(blocksize);

        minimalPhyAllocate = (long) 1024 * 1024
                * conf.getInt(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_MB,
                        YarnConfiguration.DEFAULT_RM_SCHEDULER_MINIMUM_ALLOCATION_MB);
        maximumPhyAllocate = (long) 1024 * 1024
                * conf.getInt(YarnConfiguration.RM_SCHEDULER_MAXIMUM_ALLOCATION_MB,
                        YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB);
        mrAMPhy = (long) conf.getInt(MRConfigurationNames.YARN_APP_MR_AM_RESOURCE_MB, 1536) * 1024 * 1024;

    } catch (Exception e) {
        throw new RuntimeException("Unable to analyze yarn cluster ", e);
    }

    /*
     * This is for AppMaster to query available resource in the cluster during heartbeat 
     * 
    AMRMClient<ContainerRequest> rmClient = AMRMClient.createAMRMClient();
    rmClient.init(conf);
    rmClient.start();
    AllocateResponse response = rmClient.allocate(0);
    int nodeCount = response.getNumClusterNodes();
    Resource resource = response.getAvailableResources();
    List<NodeReport> nodeUpdate = response.getUpdatedNodes();
            
    LOG.info("This is a " + nodeCount + " node cluster with totally " +
    resource.getMemory() + " memory and " + resource.getVirtualCores() + " cores");
    LOG.info(nodereport.size() + " updatedNode reports received");
    for (NodeReport node : nodeUpdate) {
       resource = node.getCapability();
       LOG.info(node.getNodeId() + " updated with " + resource.getMemory() + " memory and " + resource.getVirtualCores() + " cores");
    }*/
}

From source file:org.apache.tez.client.TestTezClient.java

License:Apache License

public void testTezClient(boolean isSession) throws Exception {
    Map<String, LocalResource> lrs = Maps.newHashMap();
    String lrName1 = "LR1";
    lrs.put(lrName1, LocalResource.newInstance(URL.newInstance("file", "localhost", 0, "/test"),
            LocalResourceType.FILE, LocalResourceVisibility.PUBLIC, 1, 1));

    TezClientForTest client = configure(lrs, isSession);

    ArgumentCaptor<ApplicationSubmissionContext> captor = ArgumentCaptor
            .forClass(ApplicationSubmissionContext.class);
    when(client.mockYarnClient.getApplicationReport(client.mockAppId).getYarnApplicationState())
            .thenReturn(YarnApplicationState.RUNNING);
    client.start();/*from  ww w  .java2 s.c  o m*/
    verify(client.mockYarnClient, times(1)).init((Configuration) any());
    verify(client.mockYarnClient, times(1)).start();
    if (isSession) {
        verify(client.mockYarnClient, times(1)).submitApplication(captor.capture());
        ApplicationSubmissionContext context = captor.getValue();
        Assert.assertEquals(3, context.getAMContainerSpec().getLocalResources().size());
        Assert.assertTrue(context.getAMContainerSpec().getLocalResources()
                .containsKey(TezConstants.TEZ_AM_LOCAL_RESOURCES_PB_FILE_NAME));
        Assert.assertTrue(context.getAMContainerSpec().getLocalResources()
                .containsKey(TezConstants.TEZ_PB_BINARY_CONF_NAME));
        Assert.assertTrue(context.getAMContainerSpec().getLocalResources().containsKey(lrName1));
    } else {
        verify(client.mockYarnClient, times(0)).submitApplication(captor.capture());
    }

    String mockLR1Name = "LR1";
    Map<String, LocalResource> lrDAG = Collections.singletonMap(mockLR1Name,
            LocalResource.newInstance(URL.newInstance("file", "localhost", 0, "/test1"), LocalResourceType.FILE,
                    LocalResourceVisibility.PUBLIC, 1, 1));
    Vertex vertex = Vertex.create("Vertex", ProcessorDescriptor.create("P"), 1, Resource.newInstance(1, 1));
    DAG dag = DAG.create("DAG").addVertex(vertex).addTaskLocalFiles(lrDAG);
    DAGClient dagClient = client.submitDAG(dag);

    Assert.assertTrue(dagClient.getExecutionContext().contains(client.mockAppId.toString()));

    if (isSession) {
        verify(client.mockYarnClient, times(1)).submitApplication(captor.capture());
        verify(client.sessionAmProxy, times(1)).submitDAG((RpcController) any(), (SubmitDAGRequestProto) any());
    } else {
        verify(client.mockYarnClient, times(1)).submitApplication(captor.capture());
        ApplicationSubmissionContext context = captor.getValue();
        Assert.assertEquals(4, context.getAMContainerSpec().getLocalResources().size());
        Assert.assertTrue(context.getAMContainerSpec().getLocalResources()
                .containsKey(TezConstants.TEZ_AM_LOCAL_RESOURCES_PB_FILE_NAME));
        Assert.assertTrue(context.getAMContainerSpec().getLocalResources()
                .containsKey(TezConstants.TEZ_PB_BINARY_CONF_NAME));
        Assert.assertTrue(context.getAMContainerSpec().getLocalResources()
                .containsKey(TezConstants.TEZ_PB_PLAN_BINARY_NAME));
        Assert.assertTrue(context.getAMContainerSpec().getLocalResources().containsKey(lrName1));
    }

    // add resources
    String lrName2 = "LR2";
    lrs.clear();
    lrs.put(lrName2, LocalResource.newInstance(URL.newInstance("file", "localhost", 0, "/test2"),
            LocalResourceType.FILE, LocalResourceVisibility.PUBLIC, 1, 1));
    client.addAppMasterLocalFiles(lrs);

    ApplicationId appId2 = ApplicationId.newInstance(0, 2);
    when(client.mockYarnClient.createApplication().getNewApplicationResponse().getApplicationId())
            .thenReturn(appId2);

    when(client.mockYarnClient.getApplicationReport(appId2).getYarnApplicationState())
            .thenReturn(YarnApplicationState.RUNNING);
    dag = DAG.create("DAG")
            .addVertex(Vertex.create("Vertex", ProcessorDescriptor.create("P"), 1, Resource.newInstance(1, 1)));
    dagClient = client.submitDAG(dag);

    if (isSession) {
        // same app master
        verify(client.mockYarnClient, times(1)).submitApplication(captor.capture());
        Assert.assertTrue(dagClient.getExecutionContext().contains(client.mockAppId.toString()));
        // additional resource is sent
        ArgumentCaptor<SubmitDAGRequestProto> captor1 = ArgumentCaptor.forClass(SubmitDAGRequestProto.class);
        verify(client.sessionAmProxy, times(2)).submitDAG((RpcController) any(), captor1.capture());
        SubmitDAGRequestProto proto = captor1.getValue();
        Assert.assertEquals(1, proto.getAdditionalAmResources().getLocalResourcesCount());
        Assert.assertEquals(lrName2, proto.getAdditionalAmResources().getLocalResources(0).getName());
    } else {
        // new app master
        Assert.assertTrue(dagClient.getExecutionContext().contains(appId2.toString()));
        verify(client.mockYarnClient, times(2)).submitApplication(captor.capture());
        // additional resource is added
        ApplicationSubmissionContext context = captor.getValue();
        Assert.assertEquals(5, context.getAMContainerSpec().getLocalResources().size());
        Assert.assertTrue(context.getAMContainerSpec().getLocalResources()
                .containsKey(TezConstants.TEZ_AM_LOCAL_RESOURCES_PB_FILE_NAME));
        Assert.assertTrue(context.getAMContainerSpec().getLocalResources()
                .containsKey(TezConstants.TEZ_PB_BINARY_CONF_NAME));
        Assert.assertTrue(context.getAMContainerSpec().getLocalResources()
                .containsKey(TezConstants.TEZ_PB_PLAN_BINARY_NAME));
        Assert.assertTrue(context.getAMContainerSpec().getLocalResources().containsKey(lrName1));
        Assert.assertTrue(context.getAMContainerSpec().getLocalResources().containsKey(lrName2));
    }

    client.stop();
    if (isSession) {
        verify(client.sessionAmProxy, times(1)).shutdownSession((RpcController) any(),
                (ShutdownSessionRequestProto) any());
    }
    verify(client.mockYarnClient, times(1)).stop();
}

From source file:org.apache.tez.client.TestTezClient.java

License:Apache License

@Test(timeout = 5000)
public void testPreWarm() throws Exception {
    TezClientForTest client = configure();
    client.start();//from ww  w . j a v a2 s  .  c  o  m

    when(client.mockYarnClient.getApplicationReport(client.mockAppId).getYarnApplicationState())
            .thenReturn(YarnApplicationState.RUNNING);

    when(client.sessionAmProxy.getAMStatus((RpcController) any(), (GetAMStatusRequestProto) any()))
            .thenReturn(GetAMStatusResponseProto.newBuilder().setStatus(TezSessionStatusProto.READY).build());

    PreWarmVertex vertex = PreWarmVertex.create("PreWarm", 1, Resource.newInstance(1, 1));
    client.preWarm(vertex);

    ArgumentCaptor<SubmitDAGRequestProto> captor1 = ArgumentCaptor.forClass(SubmitDAGRequestProto.class);
    verify(client.sessionAmProxy, times(1)).submitDAG((RpcController) any(), captor1.capture());
    SubmitDAGRequestProto proto = captor1.getValue();
    Assert.assertTrue(proto.getDAGPlan().getName().startsWith(TezConstants.TEZ_PREWARM_DAG_NAME_PREFIX));

    client.stop();
}