@Override
public void initialize() {
// Nothing to do. This class isn't currently designed to be used at the DAG API level.
UserPayload userPayload = getContext().getUserPayload();
if (userPayload == null
|| userPayload.getPayload() == null
|| userPayload.getPayload().limit() == 0) {
throw new RuntimeException(
"Could not initialize CustomShuffleEdgeManager" + " from provided user payload");
}
CustomShuffleEdgeManagerConfig config;
try {
config = CustomShuffleEdgeManagerConfig.fromUserPayload(userPayload);
} catch (InvalidProtocolBufferException e) {
throw new RuntimeException(
"Could not initialize CustomShuffleEdgeManager" + " from provided user payload", e);
}
this.numSourceTaskOutputs = config.numSourceTaskOutputs;
this.numDestinationTasks = config.numDestinationTasks;
this.basePartitionRange = config.basePartitionRange;
this.remainderRangeForLastShuffler = config.remainderRangeForLastShuffler;
this.numSourceTasks = getContext().getSourceVertexNumTasks();
Preconditions.checkState(
this.numDestinationTasks == getContext().getDestinationVertexNumTasks());
}
@Override
public void initialize(EdgeManagerContext edgeManagerContext) {
// Nothing to do. This class isn't currently designed to be used at the DAG API level.
byte[] userPayload = edgeManagerContext.getUserPayload();
if (userPayload == null || userPayload.length == 0) {
throw new RuntimeException(
"Could not initialize CustomShuffleEdgeManager" + " from provided user payload");
}
CustomShuffleEdgeManagerConfig config;
try {
config = CustomShuffleEdgeManagerConfig.fromUserPayload(userPayload);
} catch (InvalidProtocolBufferException e) {
throw new RuntimeException(
"Could not initialize CustomShuffleEdgeManager" + " from provided user payload", e);
}
this.numSourceTaskOutputs = config.numSourceTaskOutputs;
this.numDestinationTasks = config.numDestinationTasks;
this.basePartitionRange = config.basePartitionRange;
this.remainderRangeForLastShuffler = config.remainderRangeForLastShuffler;
this.numSourceTasks = config.numSourceTasks;
}
void determineParallelismAndApply() {
if (numSourceTasksCompleted == 0) {
return;
}
if (numVertexManagerEventsReceived == 0) {
return;
}
int currentParallelism = pendingTasks.size();
long expectedTotalSourceTasksOutputSize =
(numSourceTasks * completedSourceTasksOutputSize) / numVertexManagerEventsReceived;
int desiredTaskParallelism =
(int)
((expectedTotalSourceTasksOutputSize + desiredTaskInputDataSize - 1)
/ desiredTaskInputDataSize);
if (desiredTaskParallelism < minTaskParallelism) {
desiredTaskParallelism = minTaskParallelism;
}
if (desiredTaskParallelism >= currentParallelism) {
return;
}
// most shufflers will be assigned this range
int basePartitionRange = currentParallelism / desiredTaskParallelism;
if (basePartitionRange <= 1) {
// nothing to do if range is equal 1 partition. shuffler does it by default
return;
}
int numShufflersWithBaseRange = currentParallelism / basePartitionRange;
int remainderRangeForLastShuffler = currentParallelism % basePartitionRange;
int finalTaskParallelism =
(remainderRangeForLastShuffler > 0)
? (numShufflersWithBaseRange + 1)
: (numShufflersWithBaseRange);
LOG.info(
"Reduce auto parallelism for vertex: "
+ context.getVertexName()
+ " to "
+ finalTaskParallelism
+ " from "
+ pendingTasks.size()
+ " . Expected output: "
+ expectedTotalSourceTasksOutputSize
+ " based on actual output: "
+ completedSourceTasksOutputSize
+ " from "
+ numVertexManagerEventsReceived
+ " vertex manager events. "
+ " desiredTaskInputSize: "
+ desiredTaskInputDataSize);
if (finalTaskParallelism < currentParallelism) {
// final parallelism is less than actual parallelism
Map<String, EdgeManagerDescriptor> edgeManagers =
new HashMap<String, EdgeManagerDescriptor>(bipartiteSources.size());
for (String vertex : bipartiteSources.keySet()) {
// use currentParallelism for numSourceTasks to maintain original state
// for the source tasks
CustomShuffleEdgeManagerConfig edgeManagerConfig =
new CustomShuffleEdgeManagerConfig(
currentParallelism,
finalTaskParallelism,
numSourceTasks,
basePartitionRange,
((remainderRangeForLastShuffler > 0)
? remainderRangeForLastShuffler
: basePartitionRange));
EdgeManagerDescriptor edgeManagerDescriptor =
new EdgeManagerDescriptor(CustomShuffleEdgeManager.class.getName());
edgeManagerDescriptor.setUserPayload(edgeManagerConfig.toUserPayload());
edgeManagers.put(vertex, edgeManagerDescriptor);
}
context.setVertexParallelism(finalTaskParallelism, null, edgeManagers, null);
updatePendingTasks();
}
}
/**
* Compute optimal parallelism needed for the job
*
* @return true (if parallelism is determined), false otherwise
*/
@VisibleForTesting
boolean determineParallelismAndApply() {
if (numBipartiteSourceTasksCompleted == 0) {
return true;
}
if (numVertexManagerEventsReceived == 0) {
return true;
}
int currentParallelism = pendingTasks.size();
/**
* When overall completed output size is not even equal to desiredTaskInputSize, we can wait for
* some more data to be available to determine better parallelism until max.fraction is reached.
* min.fraction is just a hint to the framework and need not be honored strictly in this case.
*/
boolean canDetermineParallelismLater =
(completedSourceTasksOutputSize < desiredTaskInputDataSize)
&& (numBipartiteSourceTasksCompleted
< (totalNumBipartiteSourceTasks * slowStartMaxSrcCompletionFraction));
if (canDetermineParallelismLater) {
LOG.info(
"Defer scheduling tasks; vertex="
+ getContext().getVertexName()
+ ", totalNumBipartiteSourceTasks="
+ totalNumBipartiteSourceTasks
+ ", completedSourceTasksOutputSize="
+ completedSourceTasksOutputSize
+ ", numVertexManagerEventsReceived="
+ numVertexManagerEventsReceived
+ ", numBipartiteSourceTasksCompleted="
+ numBipartiteSourceTasksCompleted
+ ", maxThreshold="
+ (totalNumBipartiteSourceTasks * slowStartMaxSrcCompletionFraction));
return false;
}
long expectedTotalSourceTasksOutputSize =
(totalNumBipartiteSourceTasks * completedSourceTasksOutputSize)
/ numVertexManagerEventsReceived;
int desiredTaskParallelism =
(int)
((expectedTotalSourceTasksOutputSize + desiredTaskInputDataSize - 1)
/ desiredTaskInputDataSize);
if (desiredTaskParallelism < minTaskParallelism) {
desiredTaskParallelism = minTaskParallelism;
}
if (desiredTaskParallelism >= currentParallelism) {
return true;
}
// most shufflers will be assigned this range
basePartitionRange = currentParallelism / desiredTaskParallelism;
if (basePartitionRange <= 1) {
// nothing to do if range is equal 1 partition. shuffler does it by default
return true;
}
int numShufflersWithBaseRange = currentParallelism / basePartitionRange;
remainderRangeForLastShuffler = currentParallelism % basePartitionRange;
int finalTaskParallelism =
(remainderRangeForLastShuffler > 0)
? (numShufflersWithBaseRange + 1)
: (numShufflersWithBaseRange);
LOG.info(
"Reduce auto parallelism for vertex: "
+ getContext().getVertexName()
+ " to "
+ finalTaskParallelism
+ " from "
+ pendingTasks.size()
+ " . Expected output: "
+ expectedTotalSourceTasksOutputSize
+ " based on actual output: "
+ completedSourceTasksOutputSize
+ " from "
+ numVertexManagerEventsReceived
+ " vertex manager events. "
+ " desiredTaskInputSize: "
+ desiredTaskInputDataSize
+ " max slow start tasks:"
+ (totalNumBipartiteSourceTasks * slowStartMaxSrcCompletionFraction)
+ " num sources completed:"
+ numBipartiteSourceTasksCompleted);
if (finalTaskParallelism < currentParallelism) {
// final parallelism is less than actual parallelism
Map<String, EdgeProperty> edgeProperties =
new HashMap<String, EdgeProperty>(bipartiteSources);
Iterable<Map.Entry<String, SourceVertexInfo>> bipartiteItr = getBipartiteInfo();
for (Map.Entry<String, SourceVertexInfo> entry : bipartiteItr) {
String vertex = entry.getKey();
EdgeProperty oldEdgeProp = entry.getValue().edgeProperty;
// use currentParallelism for numSourceTasks to maintain original state
// for the source tasks
CustomShuffleEdgeManagerConfig edgeManagerConfig =
new CustomShuffleEdgeManagerConfig(
currentParallelism,
finalTaskParallelism,
basePartitionRange,
((remainderRangeForLastShuffler > 0)
? remainderRangeForLastShuffler
: basePartitionRange));
EdgeManagerPluginDescriptor edgeManagerDescriptor =
EdgeManagerPluginDescriptor.create(CustomShuffleEdgeManager.class.getName());
edgeManagerDescriptor.setUserPayload(edgeManagerConfig.toUserPayload());
EdgeProperty newEdgeProp =
EdgeProperty.create(
edgeManagerDescriptor,
oldEdgeProp.getDataSourceType(),
oldEdgeProp.getSchedulingType(),
oldEdgeProp.getEdgeSource(),
oldEdgeProp.getEdgeDestination());
edgeProperties.put(vertex, newEdgeProp);
}
getContext().reconfigureVertex(finalTaskParallelism, null, edgeProperties);
updatePendingTasks();
configureTargetMapping(finalTaskParallelism);
}
return true;
}
