mockInputVertices.put(mockSrcVertexId1, eProp1);
mockInputVertices.put(mockSrcVertexId2, eProp2);
mockInputVertices.put(mockSrcVertexId3, eProp3);
final VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(4);
//Check via setters
ShuffleVertexManager.ShuffleVertexManagerConfigBuilder configurer = ShuffleVertexManager
.createConfigBuilder(null);
VertexManagerPluginDescriptor pluginDesc = configurer.setAutoReduceParallelism(true)
.setDesiredTaskInputSize(1000l)
.setMinTaskParallelism(10).setSlowStartMaxSrcCompletionFraction(0.5f).build();
when(mockContext.getUserPayload()).thenReturn(pluginDesc.getUserPayload());
manager = ReflectionUtils.createClazzInstance(pluginDesc.getClassName(),
new Class[]{VertexManagerPluginContext.class}, new Object[]{mockContext});
manager.initialize();
Assert.assertTrue(manager.enableAutoParallelism == true);
Assert.assertTrue(manager.desiredTaskInputDataSize == 1000l);
Assert.assertTrue(manager.minTaskParallelism == 10);
Assert.assertTrue(manager.slowStartMinSrcCompletionFraction == 0.25f);
Assert.assertTrue(manager.slowStartMaxSrcCompletionFraction == 0.5f);
configurer = ShuffleVertexManager.createConfigBuilder(null);
pluginDesc = configurer.setAutoReduceParallelism(false).build();
when(mockContext.getUserPayload()).thenReturn(pluginDesc.getUserPayload());
manager = ReflectionUtils.createClazzInstance(pluginDesc.getClassName(),
new Class[]{VertexManagerPluginContext.class}, new Object[]{mockContext});
manager.initialize();
Assert.assertTrue(manager.enableAutoParallelism == false);
Assert.assertTrue(manager.desiredTaskInputDataSize ==
ShuffleVertexManager.TEZ_SHUFFLE_VERTEX_MANAGER_DESIRED_TASK_INPUT_SIZE_DEFAULT);
Assert.assertTrue(manager.minTaskParallelism == 1);
Assert.assertTrue(manager.slowStartMinSrcCompletionFraction ==
ShuffleVertexManager.TEZ_SHUFFLE_VERTEX_MANAGER_MIN_SRC_FRACTION_DEFAULT);
Assert.assertTrue(manager.slowStartMaxSrcCompletionFraction ==
ShuffleVertexManager.TEZ_SHUFFLE_VERTEX_MANAGER_MAX_SRC_FRACTION_DEFAULT);
// check initialization
manager = createManager(conf, mockContext, 0.1f, 0.1f);
Assert.assertTrue(manager.bipartiteSources.size() == 2);
Assert.assertTrue(manager.bipartiteSources.containsKey(mockSrcVertexId1));
Assert.assertTrue(manager.bipartiteSources.containsKey(mockSrcVertexId2));
final HashSet<Integer> scheduledTasks = new HashSet<Integer>();
doAnswer(new Answer() {
public Object answer(InvocationOnMock invocation) {
Object[] args = invocation.getArguments();
scheduledTasks.clear();
List<TaskWithLocationHint> tasks = (List<TaskWithLocationHint>)args[0];
for (TaskWithLocationHint task : tasks) {
scheduledTasks.add(task.getTaskIndex());
}
return null;
}}).when(mockContext).scheduleVertexTasks(anyList());
final Map<String, EdgeManagerPlugin> newEdgeManagers =
new HashMap<String, EdgeManagerPlugin>();
doAnswer(new Answer() {
public Object answer(InvocationOnMock invocation) {
when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(2);
newEdgeManagers.clear();
for (Entry<String, EdgeManagerPluginDescriptor> entry :
((Map<String, EdgeManagerPluginDescriptor>)invocation.getArguments()[2]).entrySet()) {
final UserPayload userPayload = entry.getValue().getUserPayload();
EdgeManagerPluginContext emContext = new EdgeManagerPluginContext() {
@Override
public UserPayload getUserPayload() {
return userPayload == null ? null : userPayload;
}
@Override
public String getSourceVertexName() {
return null;
}
@Override
public String getDestinationVertexName() {
return null;
}
@Override
public int getSourceVertexNumTasks() {
return 0;
}
@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