Examples with TaskIdentifier org.apache.tez.runtime.api.TaskIdentifier used on opensource projects

Example 1 with TaskIdentifier

use of org.apache.tez.runtime.api.TaskIdentifier in project tez by apache.

the class TestRootInputVertexManager method createTaskAttemptIdentifier.

public static TaskAttemptIdentifier createTaskAttemptIdentifier(String vName, int tId) {
    VertexIdentifier mockVertex = mock(VertexIdentifier.class);
    when(mockVertex.getName()).thenReturn(vName);
    TaskIdentifier mockTask = mock(TaskIdentifier.class);
    when(mockTask.getIdentifier()).thenReturn(tId);
    when(mockTask.getVertexIdentifier()).thenReturn(mockVertex);
    TaskAttemptIdentifier mockAttempt = mock(TaskAttemptIdentifier.class);
    when(mockAttempt.getIdentifier()).thenReturn(0);
    when(mockAttempt.getTaskIdentifier()).thenReturn(mockTask);
    return mockAttempt;
}

Example 2 with TaskIdentifier

use of org.apache.tez.runtime.api.TaskIdentifier in project tez by apache.

the class ShuffleVertexManagerBase method handleVertexManagerEvent.

private void handleVertexManagerEvent(VertexManagerEvent vmEvent) {
    // currently events from multiple attempts of the same task can be ignored because
    // their output will be the same.
    TaskIdentifier producerTask = vmEvent.getProducerAttemptIdentifier().getTaskIdentifier();
    if (!taskWithVmEvents.add(producerTask)) {
        LOG.info("Ignoring vertex manager event from: {}", producerTask);
        return;
    }
    String vName = producerTask.getVertexIdentifier().getName();
    SourceVertexInfo srcInfo = srcVertexInfo.get(vName);
    Preconditions.checkState(srcInfo != null, "Unknown vmEvent from " + producerTask);
    numVertexManagerEventsReceived++;
    long sourceTaskOutputSize = 0;
    if (vmEvent.getUserPayload() != null) {
        // save output size
        VertexManagerEventPayloadProto proto;
        try {
            proto = VertexManagerEventPayloadProto.parseFrom(ByteString.copyFrom(vmEvent.getUserPayload()));
        } catch (InvalidProtocolBufferException e) {
            throw new TezUncheckedException(e);
        }
        sourceTaskOutputSize = proto.getOutputSize();
        if (proto.hasPartitionStats()) {
            try {
                RoaringBitmap partitionStats = new RoaringBitmap();
                ByteString compressedPartitionStats = proto.getPartitionStats();
                byte[] rawData = TezCommonUtils.decompressByteStringToByteArray(compressedPartitionStats, inflater);
                NonSyncByteArrayInputStream bin = new NonSyncByteArrayInputStream(rawData);
                partitionStats.deserialize(new DataInputStream(bin));
                parsePartitionStats(srcInfo, partitionStats);
            } catch (IOException e) {
                throw new TezUncheckedException(e);
            }
        } else if (proto.hasDetailedPartitionStats()) {
            List<Integer> detailedPartitionStats = proto.getDetailedPartitionStats().getSizeInMbList();
            parseDetailedPartitionStats(srcInfo, detailedPartitionStats);
        }
        srcInfo.numVMEventsReceived++;
        srcInfo.outputSize += sourceTaskOutputSize;
        completedSourceTasksOutputSize += sourceTaskOutputSize;
    }
    if (LOG.isDebugEnabled()) {
        LOG.debug("For attempt: {} received info of output size: {}" + " vertex numEventsReceived: {} vertex output size: {}" + " total numEventsReceived: {} total output size: {}", vmEvent.getProducerAttemptIdentifier(), sourceTaskOutputSize, srcInfo.numVMEventsReceived, srcInfo.outputSize, numVertexManagerEventsReceived, completedSourceTasksOutputSize);
    }
}

Example 3 with TaskIdentifier

use of org.apache.tez.runtime.api.TaskIdentifier in project tez by apache.

the class TestShuffleVertexManagerUtils method createTaskAttemptIdentifier.

public static TaskAttemptIdentifier createTaskAttemptIdentifier(String vName, int tId) {
    VertexIdentifier mockVertex = mock(VertexIdentifier.class);
    when(mockVertex.getName()).thenReturn(vName);
    TaskIdentifier mockTask = mock(TaskIdentifier.class);
    when(mockTask.getIdentifier()).thenReturn(tId);
    when(mockTask.getVertexIdentifier()).thenReturn(mockVertex);
    TaskAttemptIdentifier mockAttempt = mock(TaskAttemptIdentifier.class);
    when(mockAttempt.getIdentifier()).thenReturn(0);
    when(mockAttempt.getTaskIdentifier()).thenReturn(mockTask);
    return mockAttempt;
}