@Override
public int getDestinationVertexNumTasks() {
return 0;
}
};
EdgeManagerPlugin edgeManager = ReflectionUtils
.createClazzInstance(entry.getValue().getClassName(),
new Class[]{EdgeManagerPluginContext.class}, new Object[]{emContext});
edgeManager.initialize();
newEdgeManagers.put(entry.getKey(), edgeManager);
}
return null;
}}).when(mockContext).setVertexParallelism(eq(2), any(VertexLocationHint.class), anyMap(), anyMap());
// source vertices have 0 tasks. immediate start of all managed tasks
when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(0);
when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(0);
when(mockContext.getVertexNumTasks(mockSrcVertexId3)).thenReturn(1);
manager.onVertexStarted(null);
Assert.assertTrue(manager.pendingTasks.isEmpty());
Assert.assertTrue(scheduledTasks.size() == 4); // all tasks scheduled
scheduledTasks.clear();
when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(2);
when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(2);
ByteBuffer payload =
VertexManagerEventPayloadProto.newBuilder().setOutputSize(5000L).build().toByteString().asReadOnlyByteBuffer();
VertexManagerEvent vmEvent = VertexManagerEvent.create("Vertex", payload);
// parallelism not change due to large data size
manager = createManager(conf, mockContext, 0.1f, 0.1f);
manager.onVertexStarted(null);
Assert.assertTrue(manager.pendingTasks.size() == 4); // no tasks scheduled
Assert.assertTrue(manager.totalNumSourceTasks == 4);
manager.onVertexManagerEventReceived(vmEvent);
manager.onSourceTaskCompleted(mockSrcVertexId1, new Integer(0));
// managedVertex tasks reduced
verify(mockContext, times(0)).setVertexParallelism(anyInt(), any(VertexLocationHint.class), anyMap(), anyMap());
Assert.assertEquals(0, manager.pendingTasks.size()); // all tasks scheduled
Assert.assertEquals(4, scheduledTasks.size());
Assert.assertEquals(1, manager.numSourceTasksCompleted);
Assert.assertEquals(5000L, manager.completedSourceTasksOutputSize);
// parallelism changed due to small data size
scheduledTasks.clear();
payload =
VertexManagerEventPayloadProto.newBuilder().setOutputSize(500L).build().toByteString().asReadOnlyByteBuffer();
vmEvent = VertexManagerEvent.create("Vertex", payload);
manager = createManager(conf, mockContext, 0.5f, 0.5f);
manager.onVertexStarted(null);
Assert.assertEquals(4, manager.pendingTasks.size()); // no tasks scheduled
Assert.assertEquals(4, manager.totalNumSourceTasks);
// task completion from non-bipartite stage does nothing
manager.onSourceTaskCompleted(mockSrcVertexId3, new Integer(0));
Assert.assertEquals(4, manager.pendingTasks.size()); // no tasks scheduled
Assert.assertEquals(4, manager.totalNumSourceTasks);
Assert.assertEquals(0, manager.numSourceTasksCompleted);
manager.onVertexManagerEventReceived(vmEvent);
manager.onSourceTaskCompleted(mockSrcVertexId1, new Integer(0));
Assert.assertEquals(4, manager.pendingTasks.size());
Assert.assertEquals(0, scheduledTasks.size()); // no tasks scheduled
Assert.assertEquals(1, manager.numSourceTasksCompleted);
Assert.assertEquals(1, manager.numVertexManagerEventsReceived);
Assert.assertEquals(500L, manager.completedSourceTasksOutputSize);
// ignore duplicate completion
manager.onSourceTaskCompleted(mockSrcVertexId1, new Integer(0));
Assert.assertEquals(4, manager.pendingTasks.size());
Assert.assertEquals(0, scheduledTasks.size()); // no tasks scheduled
Assert.assertEquals(1, manager.numSourceTasksCompleted);
Assert.assertEquals(500L, manager.completedSourceTasksOutputSize);
manager.onVertexManagerEventReceived(vmEvent);
manager.onSourceTaskCompleted(mockSrcVertexId1, new Integer(1));
// managedVertex tasks reduced
verify(mockContext).setVertexParallelism(eq(2), any(VertexLocationHint.class), anyMap(), anyMap());
Assert.assertEquals(2, newEdgeManagers.size());
// TODO improve tests for parallelism
Assert.assertEquals(0, manager.pendingTasks.size()); // all tasks scheduled
Assert.assertEquals(2, scheduledTasks.size());
Assert.assertTrue(scheduledTasks.contains(new Integer(0)));
Assert.assertTrue(scheduledTasks.contains(new Integer(1)));
Assert.assertEquals(2, manager.numSourceTasksCompleted);
Assert.assertEquals(2, manager.numVertexManagerEventsReceived);
Assert.assertEquals(1000L, manager.completedSourceTasksOutputSize);
// more completions dont cause recalculation of parallelism
manager.onSourceTaskCompleted(mockSrcVertexId2, new Integer(0));
verify(mockContext).setVertexParallelism(eq(2), any(VertexLocationHint.class), anyMap(), anyMap());
Assert.assertEquals(2, newEdgeManagers.size());
EdgeManagerPlugin edgeManager = newEdgeManagers.values().iterator().next();
Map<Integer, List<Integer>> targets = Maps.newHashMap();
DataMovementEvent dmEvent = DataMovementEvent.create(1, ByteBuffer.wrap(new byte[0]));
// 4 source task outputs - same as original number of partitions
Assert.assertEquals(4, edgeManager.getNumSourceTaskPhysicalOutputs(0));
// 4 destination task inputs - 2 source tasks + 2 merged partitions
Assert.assertEquals(4, edgeManager.getNumDestinationTaskPhysicalInputs(0));
edgeManager.routeDataMovementEventToDestination(dmEvent, 1, dmEvent.getSourceIndex(), targets);
Assert.assertEquals(1, targets.size());
Map.Entry<Integer, List<Integer>> e = targets.entrySet().iterator().next();
Assert.assertEquals(0, e.getKey().intValue());
Assert.assertEquals(1, e.getValue().size());
Assert.assertEquals(3, e.getValue().get(0).intValue());
targets.clear();
dmEvent = DataMovementEvent.create(2, ByteBuffer.wrap(new byte[0]));
edgeManager.routeDataMovementEventToDestination(dmEvent, 0, dmEvent.getSourceIndex(), targets);
Assert.assertEquals(1, targets.size());
e = targets.entrySet().iterator().next();
Assert.assertEquals(1, e.getKey().intValue());
Assert.assertEquals(1, e.getValue().size());
Assert.assertEquals(0, e.getValue().get(0).intValue());
targets.clear();
edgeManager.routeInputSourceTaskFailedEventToDestination(2, targets);
Assert.assertEquals(2, targets.size());
for (Map.Entry<Integer, List<Integer>> entry : targets.entrySet()) {
Assert.assertTrue(entry.getKey().intValue() == 0 || entry.getKey().intValue() == 1);
Assert.assertEquals(2, entry.getValue().size());
Assert.assertEquals(4, entry.getValue().get(0).intValue());