Examples with PreemptionMessage org.apache.hadoop.yarn.api.records.PreemptionMessage used on opensource projects

Example 6 with PreemptionMessage

use of org.apache.hadoop.yarn.api.records.PreemptionMessage in project hadoop by apache.

the class TestKillAMPreemptionPolicy method testKillAMPreemptPolicy.

@SuppressWarnings("unchecked")
@Test
public void testKillAMPreemptPolicy() {
    ApplicationId appId = ApplicationId.newInstance(123456789, 1);
    ContainerId container = ContainerId.newContainerId(ApplicationAttemptId.newInstance(appId, 1), 1);
    AMPreemptionPolicy.Context mPctxt = mock(AMPreemptionPolicy.Context.class);
    when(mPctxt.getTaskAttempt(any(ContainerId.class))).thenReturn(MRBuilderUtils.newTaskAttemptId(MRBuilderUtils.newTaskId(MRBuilderUtils.newJobId(appId, 1), 1, TaskType.MAP), 0));
    List<Container> p = new ArrayList<Container>();
    p.add(Container.newInstance(container, null, null, null, null, null));
    when(mPctxt.getContainers(any(TaskType.class))).thenReturn(p);
    KillAMPreemptionPolicy policy = new KillAMPreemptionPolicy();
    // strictContract is null & contract is null
    RunningAppContext mActxt = getRunningAppContext();
    policy.init(mActxt);
    PreemptionMessage pM = getPreemptionMessage(false, false, container);
    policy.preempt(mPctxt, pM);
    verify(mActxt.getEventHandler(), times(0)).handle(any(TaskAttemptEvent.class));
    verify(mActxt.getEventHandler(), times(0)).handle(any(JobCounterUpdateEvent.class));
    // strictContract is not null & contract is null
    mActxt = getRunningAppContext();
    policy.init(mActxt);
    pM = getPreemptionMessage(true, false, container);
    policy.preempt(mPctxt, pM);
    verify(mActxt.getEventHandler(), times(2)).handle(any(TaskAttemptEvent.class));
    verify(mActxt.getEventHandler(), times(2)).handle(any(JobCounterUpdateEvent.class));
    // strictContract is null & contract is not null
    mActxt = getRunningAppContext();
    policy.init(mActxt);
    pM = getPreemptionMessage(false, true, container);
    policy.preempt(mPctxt, pM);
    verify(mActxt.getEventHandler(), times(2)).handle(any(TaskAttemptEvent.class));
    verify(mActxt.getEventHandler(), times(2)).handle(any(JobCounterUpdateEvent.class));
    // strictContract is not null & contract is not null
    mActxt = getRunningAppContext();
    policy.init(mActxt);
    pM = getPreemptionMessage(true, true, container);
    policy.preempt(mPctxt, pM);
    verify(mActxt.getEventHandler(), times(4)).handle(any(TaskAttemptEvent.class));
    verify(mActxt.getEventHandler(), times(4)).handle(any(JobCounterUpdateEvent.class));
}

Example 7 with PreemptionMessage

use of org.apache.hadoop.yarn.api.records.PreemptionMessage in project hadoop by apache.

the class TestKillAMPreemptionPolicy method getPreemptionMessage.

private PreemptionMessage getPreemptionMessage(boolean strictContract, boolean contract, final ContainerId container) {
    PreemptionMessage preemptionMessage = recordFactory.newRecordInstance(PreemptionMessage.class);
    Set<PreemptionContainer> cntrs = new HashSet<PreemptionContainer>();
    PreemptionContainer preemptContainer = recordFactory.newRecordInstance(PreemptionContainer.class);
    preemptContainer.setId(container);
    cntrs.add(preemptContainer);
    if (strictContract) {
        StrictPreemptionContract set = recordFactory.newRecordInstance(StrictPreemptionContract.class);
        set.setContainers(cntrs);
        preemptionMessage.setStrictContract(set);
    }
    if (contract) {
        PreemptionContract preemptContract = recordFactory.newRecordInstance(PreemptionContract.class);
        preemptContract.setContainers(cntrs);
        preemptionMessage.setContract(preemptContract);
    }
    return preemptionMessage;
}