Ejemplo n.º 1
0
 @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());
 }
Ejemplo n.º 2
0
 public static CustomShuffleEdgeManagerConfig fromUserPayload(UserPayload payload)
     throws InvalidProtocolBufferException {
   ShuffleEdgeManagerConfigPayloadProto proto =
       ShuffleEdgeManagerConfigPayloadProto.parseFrom(ByteString.copyFrom(payload.getPayload()));
   return new CustomShuffleEdgeManagerConfig(
       proto.getNumSourceTaskOutputs(),
       proto.getNumDestinationTasks(),
       proto.getBasePartitionRange(),
       proto.getRemainderRangeForLastShuffler());
 }
  @Test(timeout = 5000)
  public void testGetBytePayload() throws IOException {
    int numBuckets = 10;
    VertexManagerPluginContext context = mock(VertexManagerPluginContext.class);
    CustomVertexConfiguration vertexConf =
        new CustomVertexConfiguration(numBuckets, TezWork.VertexType.INITIALIZED_EDGES);
    DataOutputBuffer dob = new DataOutputBuffer();
    vertexConf.write(dob);
    UserPayload payload = UserPayload.create(ByteBuffer.wrap(dob.getData()));
    when(context.getUserPayload()).thenReturn(payload);

    CustomPartitionVertex vm = new CustomPartitionVertex(context);
    vm.initialize();

    // prepare empty routing table
    Multimap<Integer, Integer> routingTable = HashMultimap.<Integer, Integer>create();
    payload = vm.getBytePayload(routingTable);
    // get conf from user payload
    CustomEdgeConfiguration edgeConf = new CustomEdgeConfiguration();
    DataInputByteBuffer dibb = new DataInputByteBuffer();
    dibb.reset(payload.getPayload());
    edgeConf.readFields(dibb);
    assertEquals(numBuckets, edgeConf.getNumBuckets());
  }