Examples with EdgeManagerPluginContext org.apache.tez.dag.api.EdgeManagerPluginContext used on opensource projects

Example 1 with EdgeManagerPluginContext

use of org.apache.tez.dag.api.EdgeManagerPluginContext in project tez by apache.

the class TestEdge method testOneToOneEdgeManagerODR.

@Test(timeout = 5000)
public void testOneToOneEdgeManagerODR() {
    EdgeManagerPluginContext mockContext = mock(EdgeManagerPluginContext.class);
    when(mockContext.getSourceVertexName()).thenReturn("Source");
    when(mockContext.getDestinationVertexName()).thenReturn("Destination");
    when(mockContext.getSourceVertexNumTasks()).thenReturn(3);
    OneToOneEdgeManagerOnDemand manager = new OneToOneEdgeManagerOnDemand(mockContext);
    manager.initialize();
    Map<Integer, List<Integer>> destinationTaskAndInputIndices = Maps.newHashMap();
    DataMovementEvent event = DataMovementEvent.create(1, null);
    // fail when source and destination are inconsistent
    when(mockContext.getDestinationVertexNumTasks()).thenReturn(4);
    try {
        manager.routeDataMovementEventToDestination(event, 1, 1, destinationTaskAndInputIndices);
        Assert.fail();
    } catch (IllegalStateException e) {
        Assert.assertTrue(e.getMessage().contains("1-1 source and destination task counts must match"));
    }
    // now make it consistent
    when(mockContext.getDestinationVertexNumTasks()).thenReturn(3);
    manager.routeDataMovementEventToDestination(event, 1, 1, destinationTaskAndInputIndices);
    Assert.assertEquals(1, destinationTaskAndInputIndices.size());
    Assert.assertEquals(1, destinationTaskAndInputIndices.entrySet().iterator().next().getKey().intValue());
    Assert.assertEquals(0, destinationTaskAndInputIndices.entrySet().iterator().next().getValue().get(0).intValue());
}

Example 2 with EdgeManagerPluginContext

use of org.apache.tez.dag.api.EdgeManagerPluginContext in project tez by apache.

the class TestEdge method testScatterGatherManager.

@Test(timeout = 5000)
public void testScatterGatherManager() {
    EdgeManagerPluginContext mockContext = mock(EdgeManagerPluginContext.class);
    when(mockContext.getSourceVertexName()).thenReturn("Source");
    when(mockContext.getDestinationVertexName()).thenReturn("Destination");
    ScatterGatherEdgeManager manager = new ScatterGatherEdgeManager(mockContext);
    manager.initialize();
    when(mockContext.getDestinationVertexNumTasks()).thenReturn(-1);
    try {
        manager.getNumSourceTaskPhysicalOutputs(0);
        Assert.fail();
    } catch (IllegalArgumentException e) {
        e.printStackTrace();
        Assert.assertTrue(e.getMessage().contains("ScatteGather edge manager must have destination vertex task parallelism specified"));
    }
    when(mockContext.getDestinationVertexNumTasks()).thenReturn(0);
    manager.getNumSourceTaskPhysicalOutputs(0);
}

Example 3 with EdgeManagerPluginContext

use of org.apache.tez.dag.api.EdgeManagerPluginContext in project tez by apache.

the class TestEdge method testOneToOneEdgeManager.

@Test(timeout = 5000)
public void testOneToOneEdgeManager() {
    EdgeManagerPluginContext mockContext = mock(EdgeManagerPluginContext.class);
    when(mockContext.getSourceVertexName()).thenReturn("Source");
    when(mockContext.getDestinationVertexName()).thenReturn("Destination");
    when(mockContext.getSourceVertexNumTasks()).thenReturn(3);
    OneToOneEdgeManager manager = new OneToOneEdgeManager(mockContext);
    manager.initialize();
    Map<Integer, List<Integer>> destinationTaskAndInputIndices = Maps.newHashMap();
    DataMovementEvent event = DataMovementEvent.create(1, null);
    // fail when source and destination are inconsistent
    when(mockContext.getDestinationVertexNumTasks()).thenReturn(4);
    try {
        manager.routeDataMovementEventToDestination(event, 1, 1, destinationTaskAndInputIndices);
        Assert.fail();
    } catch (IllegalStateException e) {
        Assert.assertTrue(e.getMessage().contains("1-1 source and destination task counts must match"));
    }
    // now make it consistent
    when(mockContext.getDestinationVertexNumTasks()).thenReturn(3);
    manager.routeDataMovementEventToDestination(event, 1, 1, destinationTaskAndInputIndices);
    Assert.assertEquals(1, destinationTaskAndInputIndices.size());
    Assert.assertEquals(1, destinationTaskAndInputIndices.entrySet().iterator().next().getKey().intValue());
    Assert.assertEquals(0, destinationTaskAndInputIndices.entrySet().iterator().next().getValue().get(0).intValue());
}

Example 4 with EdgeManagerPluginContext

use of org.apache.tez.dag.api.EdgeManagerPluginContext in project tez by apache.

the class TestCartesianProductEdgeManager method testInitialize.

@Test(timeout = 5000)
public void testInitialize() throws Exception {
    EdgeManagerPluginContext context = mock(EdgeManagerPluginContext.class);
    when(context.getSourceVertexName()).thenReturn("v0");
    CartesianProductEdgeManager edgeManager = new CartesianProductEdgeManager(context);
    // partitioned case
    CartesianProductConfigProto.Builder builder = CartesianProductConfigProto.newBuilder();
    builder.setIsPartitioned(true).addAllSources(Arrays.asList("v0", "v1")).addAllNumPartitions(Ints.asList(2, 3)).setMaxParallelism(100).setMinOpsPerWorker(1);
    UserPayload payload = UserPayload.create(ByteBuffer.wrap(builder.build().toByteArray()));
    when(context.getUserPayload()).thenReturn(payload);
    edgeManager.initialize();
    assertTrue(edgeManager.getEdgeManagerReal() instanceof CartesianProductEdgeManagerPartitioned);
    // unpartitioned case
    builder.clear();
    builder.setIsPartitioned(false).addAllSources(Arrays.asList("v0", "v1")).addAllNumChunks(Ints.asList(2, 3)).setMaxParallelism(100).setMinOpsPerWorker(1);
    payload = UserPayload.create(ByteBuffer.wrap(builder.build().toByteArray()));
    when(context.getUserPayload()).thenReturn(payload);
    when(context.getSourceVertexNumTasks()).thenReturn(2);
    edgeManager.initialize();
    assertTrue(edgeManager.getEdgeManagerReal() instanceof FairCartesianProductEdgeManager);
}