Examples with TaskAttemptIdentifierImpl org.apache.tez.dag.records.TaskAttemptIdentifierImpl used on opensource projects

Example 1 with TaskAttemptIdentifierImpl

use of org.apache.tez.dag.records.TaskAttemptIdentifierImpl in project tez by apache.

the class TestCartesianProductVertexManagerPartitioned method setupWithConfig.

private void setupWithConfig(CartesianProductConfigProto config) throws TezReflectionException {
    MockitoAnnotations.initMocks(this);
    context = mock(VertexManagerPluginContext.class);
    when(context.getVertexName()).thenReturn("cp");
    when(context.getVertexNumTasks("cp")).thenReturn(-1);
    vertexManager = new CartesianProductVertexManagerPartitioned(context);
    Map<String, EdgeProperty> edgePropertyMap = new HashMap<>();
    edgePropertyMap.put("v0", EdgeProperty.create(EdgeManagerPluginDescriptor.create(CartesianProductEdgeManager.class.getName()), null, null, null, null));
    edgePropertyMap.put("v1", EdgeProperty.create(EdgeManagerPluginDescriptor.create(CartesianProductEdgeManager.class.getName()), null, null, null, null));
    edgePropertyMap.put("v2", EdgeProperty.create(BROADCAST, null, null, null, null));
    when(context.getInputVertexEdgeProperties()).thenReturn(edgePropertyMap);
    when(context.getVertexNumTasks(eq("v0"))).thenReturn(4);
    when(context.getVertexNumTasks(eq("v1"))).thenReturn(4);
    when(context.getVertexNumTasks(eq("v2"))).thenReturn(4);
    vertexManager.initialize(config);
    allCompletions = new ArrayList<>();
    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 4; j++) {
            allCompletions.add(new TaskAttemptIdentifierImpl("dag", "v" + i, TezTaskAttemptID.getInstance(TezTaskID.getInstance(TezVertexID.getInstance(TezDAGID.getInstance("0", 0, 0), i), j), 0)));
        }
    }
}

Example 2 with TaskAttemptIdentifierImpl

use of org.apache.tez.dag.records.TaskAttemptIdentifierImpl in project tez by apache.

the class TestShuffleVertexManagerBase method testVMEventFirst.

// Test vmEvent and vertexStatusUpdate before started.
@Test(timeout = 5000)
public void testVMEventFirst() throws IOException {
    Configuration conf = new Configuration();
    ShuffleVertexManagerBase manager;
    final String mockSrcVertexId1 = "Vertex1";
    final String mockSrcVertexId2 = "Vertex2";
    final String mockSrcVertexId3 = "Vertex3";
    final String mockManagedVertexId = "Vertex4";
    final List<Integer> scheduledTasks = Lists.newLinkedList();
    final VertexManagerPluginContext mockContext = createVertexManagerContext(mockSrcVertexId1, 2, mockSrcVertexId2, 2, mockSrcVertexId3, 2, mockManagedVertexId, 4, scheduledTasks, null);
    VertexManagerEvent vmEvent = getVertexManagerEvent(null, 1L, "Vertex");
    manager = createManager(conf, mockContext, 0.01f, 0.75f);
    // no tasks scheduled
    Assert.assertEquals(4, manager.pendingTasks.size());
    Assert.assertEquals(0, manager.numBipartiteSourceTasksCompleted);
    TezTaskAttemptID taId1 = TezTaskAttemptID.fromString("attempt_1436907267600_195589_1_00_000000_0");
    vmEvent.setProducerAttemptIdentifier(new TaskAttemptIdentifierImpl("dag", mockSrcVertexId1, taId1));
    manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
    manager.onVertexManagerEventReceived(vmEvent);
    // nothing happens
    Assert.assertEquals(0, manager.numVertexManagerEventsReceived);
    // now the processing happens
    manager.onVertexStarted(emptyCompletions);
    Assert.assertEquals(1, manager.numVertexManagerEventsReceived);
}

Example 3 with TaskAttemptIdentifierImpl

use of org.apache.tez.dag.records.TaskAttemptIdentifierImpl in project tez by apache.

the class TestShuffleVertexManagerBase method testPartitionStats.

// Test partition stats.
@Test(timeout = 5000)
public void testPartitionStats() throws IOException {
    Configuration conf = new Configuration();
    ShuffleVertexManagerBase manager;
    final String mockSrcVertexId1 = "Vertex1";
    final String mockSrcVertexId2 = "Vertex2";
    final String mockSrcVertexId3 = "Vertex3";
    final String mockManagedVertexId = "Vertex4";
    final List<Integer> scheduledTasks = Lists.newLinkedList();
    final VertexManagerPluginContext mockContext = createVertexManagerContext(mockSrcVertexId1, 2, mockSrcVertexId2, 2, mockSrcVertexId3, 2, mockManagedVertexId, 4, scheduledTasks, null);
    // {5,9,12,18} in bitmap
    final long MB = 1024l * 1024l;
    long[] sizes = new long[] { (0l), (1 * MB), (964 * MB), (48 * MB) };
    VertexManagerEvent vmEvent = getVertexManagerEvent(sizes, 0, "Vertex", false);
    manager = createManager(conf, mockContext, 0.01f, 0.75f);
    manager.onVertexStarted(emptyCompletions);
    // no tasks scheduled
    Assert.assertEquals(4, manager.pendingTasks.size());
    Assert.assertEquals(0, manager.numBipartiteSourceTasksCompleted);
    TezTaskAttemptID taId1 = TezTaskAttemptID.fromString("attempt_1436907267600_195589_1_00_000000_0");
    vmEvent.setProducerAttemptIdentifier(new TaskAttemptIdentifierImpl("dag", mockSrcVertexId1, taId1));
    manager.onVertexManagerEventReceived(vmEvent);
    Assert.assertEquals(1, manager.numVertexManagerEventsReceived);
    // 0 MB bucket
    Assert.assertEquals(0, manager.getCurrentlyKnownStatsAtIndex(0));
    // 1 MB bucket
    Assert.assertEquals(1, manager.getCurrentlyKnownStatsAtIndex(1));
    // 100 MB bucket
    Assert.assertEquals(100, manager.getCurrentlyKnownStatsAtIndex(2));
    // 10 MB bucket
    Assert.assertEquals(10, manager.getCurrentlyKnownStatsAtIndex(3));
    // sending again from a different version of the same task has not impact
    TezTaskAttemptID taId2 = TezTaskAttemptID.fromString("attempt_1436907267600_195589_1_00_000000_1");
    vmEvent.setProducerAttemptIdentifier(new TaskAttemptIdentifierImpl("dag", mockSrcVertexId1, taId2));
    manager.onVertexManagerEventReceived(vmEvent);
    Assert.assertEquals(1, manager.numVertexManagerEventsReceived);
    // 0 MB bucket
    Assert.assertEquals(0, manager.getCurrentlyKnownStatsAtIndex(0));
    // 1 MB bucket
    Assert.assertEquals(1, manager.getCurrentlyKnownStatsAtIndex(1));
    // 100 MB bucket
    Assert.assertEquals(100, manager.getCurrentlyKnownStatsAtIndex(2));
    // 10 MB bucket
    Assert.assertEquals(10, manager.getCurrentlyKnownStatsAtIndex(3));
    // Testing for detailed partition stats
    vmEvent = getVertexManagerEvent(sizes, 0, "Vertex", true);
    manager = createManager(conf, mockContext, 0.01f, 0.75f);
    manager.onVertexStarted(emptyCompletions);
    // no tasks scheduled
    Assert.assertEquals(4, manager.pendingTasks.size());
    Assert.assertEquals(0, manager.numBipartiteSourceTasksCompleted);
    taId1 = TezTaskAttemptID.fromString("attempt_1436907267600_195589_1_00_000000_0");
    vmEvent.setProducerAttemptIdentifier(new TaskAttemptIdentifierImpl("dag", mockSrcVertexId1, taId1));
    manager.onVertexManagerEventReceived(vmEvent);
    Assert.assertEquals(1, manager.numVertexManagerEventsReceived);
    Assert.assertEquals(0, manager.getCurrentlyKnownStatsAtIndex(0));
    Assert.assertEquals(1, manager.getCurrentlyKnownStatsAtIndex(1));
    Assert.assertEquals(964, manager.getCurrentlyKnownStatsAtIndex(2));
    Assert.assertEquals(48, manager.getCurrentlyKnownStatsAtIndex(3));
    // sending again from a different version of the same task has not impact
    taId2 = TezTaskAttemptID.fromString("attempt_1436907267600_195589_1_00_000000_1");
    vmEvent.setProducerAttemptIdentifier(new TaskAttemptIdentifierImpl("dag", mockSrcVertexId1, taId2));
    manager.onVertexManagerEventReceived(vmEvent);
    Assert.assertEquals(1, manager.numVertexManagerEventsReceived);
    Assert.assertEquals(0, manager.getCurrentlyKnownStatsAtIndex(0));
    Assert.assertEquals(1, manager.getCurrentlyKnownStatsAtIndex(1));
    Assert.assertEquals(964, manager.getCurrentlyKnownStatsAtIndex(2));
    Assert.assertEquals(48, manager.getCurrentlyKnownStatsAtIndex(3));
}

Example 4 with TaskAttemptIdentifierImpl

use of org.apache.tez.dag.records.TaskAttemptIdentifierImpl in project tez by apache.

the class TestShuffleVertexManagerUtils method getVertexManagerEvent.

VertexManagerEvent getVertexManagerEvent(long[] partitionSizes, long uncompressedTotalSize, String vertexName, boolean reportDetailedStats) throws IOException {
    ByteBuffer payload;
    long totalSize = 0;
    // Use partition sizes to compute the total size.
    if (partitionSizes != null) {
        totalSize = estimatedUncompressedSum(partitionSizes);
    } else {
        totalSize = uncompressedTotalSize;
    }
    if (partitionSizes != null) {
        RoaringBitmap partitionStats = ShuffleUtils.getPartitionStatsForPhysicalOutput(partitionSizes);
        DataOutputBuffer dout = new DataOutputBuffer();
        partitionStats.serialize(dout);
        ByteString partitionStatsBytes = TezCommonUtils.compressByteArrayToByteString(dout.getData());
        if (reportDetailedStats) {
            payload = VertexManagerEventPayloadProto.newBuilder().setOutputSize(totalSize).setDetailedPartitionStats(ShuffleUtils.getDetailedPartitionStatsForPhysicalOutput(partitionSizes)).build().toByteString().asReadOnlyByteBuffer();
        } else {
            payload = VertexManagerEventPayloadProto.newBuilder().setOutputSize(totalSize).setPartitionStats(partitionStatsBytes).build().toByteString().asReadOnlyByteBuffer();
        }
    } else {
        payload = VertexManagerEventPayloadProto.newBuilder().setOutputSize(totalSize).build().toByteString().asReadOnlyByteBuffer();
    }
    TaskAttemptIdentifierImpl taId = new TaskAttemptIdentifierImpl("dag", vertexName, TezTaskAttemptID.getInstance(TezTaskID.getInstance(vertexId, taskId++), 0));
    VertexManagerEvent vmEvent = VertexManagerEvent.create(vertexName, payload);
    vmEvent.setProducerAttemptIdentifier(taId);
    return vmEvent;
}