org.apache.tez.dag.api.InputDescriptor Java Examples

@Test(timeout = 5000)
public void testScalingNoProcessor2() throws TezException {
  // Real world values
  MemoryDistributor dist = new MemoryDistributor(2, 0, conf);
  
  dist.setJvmMemory(209715200l);

  // First request
  MemoryUpdateCallbackForTest e1Callback = new MemoryUpdateCallbackForTest();
  InputContext e1InputContext1 = createTestInputContext();
  InputDescriptor e1InDesc1 = createTestInputDescriptor();
  dist.requestMemory(104857600l, e1Callback, e1InputContext1, e1InDesc1);
  
  // Second request
  MemoryUpdateCallbackForTest e2Callback = new MemoryUpdateCallbackForTest();
  InputContext e2InputContext2 = createTestInputContext();
  InputDescriptor e2InDesc2 = createTestInputDescriptor();
  dist.requestMemory(157286400l, e2Callback, e2InputContext2, e2InDesc2);
  
  dist.makeInitialAllocations();

  assertEquals(58720256l, e1Callback.assigned);
  assertEquals(88080384l, e2Callback.assigned);
}
 

@Test(timeout = 5000)
public void testInvalidSourceTaskIndex() throws Exception {
  EventHandler mockEventHandler = mock(EventHandler.class);
  Edge edge = new Edge(EdgeProperty.create(
      EdgeManagerPluginDescriptor.create(CustomEdgeManagerWithInvalidReturnValue.class.getName())
        .setUserPayload(new CustomEdgeManagerWithInvalidReturnValue.EdgeManagerConfig(1,1,1,-1).toUserPayload()),
      DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(""),
      InputDescriptor.create("")), mockEventHandler, new TezConfiguration());
  TezVertexID v1Id = createVertexID(1);
  TezVertexID v2Id = createVertexID(2);
  edge.setSourceVertex(mockVertex("v1", v1Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.setDestinationVertex(mockVertex("v2", v2Id, new LinkedHashMap<TezTaskID, Task>()));
  edge.initialize();
  try {
    TezEvent ireEvent = new TezEvent(InputReadErrorEvent.create("diag", 0, 1),
        new EventMetaData(EventProducerConsumerType.INPUT, "v2", "v1",
            TezTaskAttemptID.getInstance(TezTaskID.getInstance(v2Id, 1), 1)));
    edge.sendTezEventToSourceTasks(ireEvent);
    Assert.fail();
  } catch (AMUserCodeException e) {
    e.printStackTrace();
    assertTrue(e.getCause().getMessage().contains("SourceTaskIndex should not be negative"));
  }
}
 

@Override
public synchronized List<InputSpec> getInputSpecList(int taskIndex) {
  inputSpecList = new ArrayList<InputSpec>(this.getInputVerticesCount()
      + (rootInputDescriptors == null ? 0 : rootInputDescriptors.size()));
  if (rootInputDescriptors != null) {
    for (Entry<String, RootInputLeafOutputDescriptor<InputDescriptor>> rootInputDescriptorEntry : rootInputDescriptors
        .entrySet()) {
      inputSpecList.add(new InputSpec(rootInputDescriptorEntry.getKey(),
          rootInputDescriptorEntry.getValue().getDescriptor(), rootInputSpecs.get(
              rootInputDescriptorEntry.getKey()).getNumPhysicalInputsForWorkUnit(taskIndex)));
    }
  }
  for (Entry<Vertex, Edge> entry : this.getInputVertices().entrySet()) {
    InputSpec inputSpec = entry.getValue().getDestinationSpec(taskIndex);
    if (LOG.isDebugEnabled()) {
      LOG.debug("For vertex : " + this.getName()
          + ", Using InputSpec : " + inputSpec);
    }
    // TODO DAGAM This should be based on the edge type.
    inputSpecList.add(inputSpec);
  }
  return inputSpecList;
}
 

private GroupInputEdge newGroupInputEdge(TezOperator fromOp,
        TezOperator toOp, VertexGroup from, Vertex to) throws IOException {

    EdgeProperty edgeProperty = newEdge(fromOp, toOp);

    String groupInputClass = ConcatenatedMergedKeyValueInput.class.getName();

    // In case of SCATTER_GATHER and UnorderedKVInput it will still be
    // ConcatenatedMergedKeyValueInput
    if(edgeProperty.getDataMovementType().equals(DataMovementType.SCATTER_GATHER)
            && edgeProperty.getEdgeDestination().getClassName().equals(OrderedGroupedKVInput.class.getName())) {
        groupInputClass = OrderedGroupedMergedKVInput.class.getName();
    }

    return GroupInputEdge.create(from, to, edgeProperty,
            InputDescriptor.create(groupInputClass).setUserPayload(edgeProperty.getEdgeDestination().getUserPayload())
                .setHistoryText(edgeProperty.getEdgeDestination().getHistoryText()));
}
 

@Override
public void setAdditionalInputs(List<RootInputLeafOutputProto> inputs) {
  LOG.info("Setting " + inputs.size() + " additional inputs for vertex" + this.logIdentifier);
  this.rootInputDescriptors = Maps.newHashMapWithExpectedSize(inputs.size());
  for (RootInputLeafOutputProto input : inputs) {
    addIO(input.getName());
    InputDescriptor id = DagTypeConverters
        .convertInputDescriptorFromDAGPlan(input.getIODescriptor());

    this.rootInputDescriptors
        .put(
            input.getName(),
            new RootInputLeafOutput<InputDescriptor, InputInitializerDescriptor>(
                input.getName(), id,
                input.hasControllerDescriptor() ? DagTypeConverters
                    .convertInputInitializerDescriptorFromDAGPlan(input
                        .getControllerDescriptor()) : null));
    this.rootInputSpecs.put(input.getName(), DEFAULT_ROOT_INPUT_SPECS);
  }
}
 

private AMSchedulerEventTALaunchRequest createLaunchRequestEvent(
  TezTaskAttemptID taID, TaskAttempt ta, Resource capability,
  String[] hosts, String[] racks, Priority priority,
  ContainerContext containerContext) {
  TaskLocationHint locationHint = null;
  if (hosts != null || racks != null) {
    Set<String> hostsSet = Sets.newHashSet(hosts);
    Set<String> racksSet = Sets.newHashSet(racks);
    locationHint = TaskLocationHint.createTaskLocationHint(hostsSet, racksSet);
  }
  AMSchedulerEventTALaunchRequest lr = new AMSchedulerEventTALaunchRequest(
    taID, capability, new TaskSpec(taID, "dagName", "vertexName", -1,
      ProcessorDescriptor.create("processorClassName"),
    Collections.singletonList(new InputSpec("vertexName",
        InputDescriptor.create("inputClassName"), 1)),
    Collections.singletonList(new OutputSpec("vertexName",
        OutputDescriptor.create("outputClassName"), 1)), null, null), ta, locationHint,
    priority.getPriority(), containerContext, 0, 0, 0);
  return lr;
}
 

private DAG createDAG(String dagName, boolean uv12CommitFail, boolean v3CommitFail) {
  DAG dag = DAG.create(dagName);
  Vertex v1 = Vertex.create("v1", ProcessorDescriptor.create("Proc"), 1);
  Vertex v2 = Vertex.create("v2", ProcessorDescriptor.create("Proc"), 1);
  Vertex v3 = Vertex.create("v3", ProcessorDescriptor.create("Proc"), 1);
  VertexGroup uv12 = dag.createVertexGroup("uv12", v1, v2);
  DataSinkDescriptor uv12DataSink = DataSinkDescriptor.create(
      OutputDescriptor.create("dummy output"), createOutputCommitterDesc(uv12CommitFail), null);
  uv12.addDataSink("uv12Out", uv12DataSink);
  DataSinkDescriptor v3DataSink = DataSinkDescriptor.create(
      OutputDescriptor.create("dummy output"), createOutputCommitterDesc(v3CommitFail), null);
  v3.addDataSink("v3Out", v3DataSink);

  GroupInputEdge e1 = GroupInputEdge.create(uv12, v3, EdgeProperty.create(
      DataMovementType.SCATTER_GATHER, DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create("dummy output class"),
      InputDescriptor.create("dummy input class")), InputDescriptor
      .create("merge.class"));
  dag.addVertex(v1)
    .addVertex(v2)
    .addVertex(v3)
    .addEdge(e1);
  return dag;
}
 

@Private
private static DataSourceDescriptor configureMRInputWithLegacySplitsGenerated(Configuration conf,
                                                                              boolean useLegacyInput) {
  InputDescriptor inputDescriptor;

  try {
    inputDescriptor = InputDescriptor.create(useLegacyInput ? MRInputLegacy.class
        .getName() : MRInput.class.getName())
        .setUserPayload(MRInputHelpersInternal.createMRInputPayload(conf, null));
  } catch (IOException e) {
    throw new TezUncheckedException(e);
  }

  DataSourceDescriptor dsd = DataSourceDescriptor.create(inputDescriptor, null, null);
  if (conf.getBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_CONVERT_USER_PAYLOAD_TO_HISTORY_TEXT,
      TezRuntimeConfiguration.TEZ_RUNTIME_CONVERT_USER_PAYLOAD_TO_HISTORY_TEXT_DEFAULT)) {
    dsd.getInputDescriptor().setHistoryText(TezUtils.convertToHistoryText(conf));
  }

  return dsd;
}
 

private DAG createDAG(TezConfiguration tezConf) throws IOException {
  InputDescriptor inputDescriptor = InputDescriptor.create(FakeInput.class.getName());
  InputInitializerDescriptor inputInitializerDescriptor =
    InputInitializerDescriptor.create(FakeInputInitializer.class.getName());
  DataSourceDescriptor dataSourceDescriptor =
    DataSourceDescriptor.create(inputDescriptor, inputInitializerDescriptor, null);

  Vertex v1 = Vertex.create(VERTEX1, ProcessorDescriptor.create(TokenProcessor.class.getName()));
  v1.addDataSource(INPUT, dataSourceDescriptor);
  Vertex v2 = Vertex.create(VERTEX2, ProcessorDescriptor.create(TokenProcessor.class.getName()));
  v2.addDataSource(INPUT, dataSourceDescriptor);

  OutputDescriptor outputDescriptor = OutputDescriptor.create(FakeOutput.class.getName());
  OutputCommitterDescriptor outputCommitterDescriptor =
    OutputCommitterDescriptor.create(FakeOutputCommitter.class.getName());
  DataSinkDescriptor dataSinkDescriptor =
    DataSinkDescriptor.create(outputDescriptor, outputCommitterDescriptor, null);

  CartesianProductConfig cartesianProductConfig =
    new CartesianProductConfig(Arrays.asList(sourceVertices));
  UserPayload userPayload = cartesianProductConfig.toUserPayload(tezConf);

  Vertex v3 = Vertex.create(VERTEX3, ProcessorDescriptor.create(JoinProcessor.class.getName()));
  v3.addDataSink(OUTPUT, dataSinkDescriptor);
  v3.setVertexManagerPlugin(
    VertexManagerPluginDescriptor.create(CartesianProductVertexManager.class.getName())
                                 .setUserPayload(userPayload));

  EdgeManagerPluginDescriptor edgeManagerDescriptor =
    EdgeManagerPluginDescriptor.create(CartesianProductEdgeManager.class.getName());
  edgeManagerDescriptor.setUserPayload(userPayload);
  UnorderedPartitionedKVEdgeConfig edgeConf =
    UnorderedPartitionedKVEdgeConfig.newBuilder(Text.class.getName(), IntWritable.class.getName(),
      RoundRobinPartitioner.class.getName()).build();
  EdgeProperty edgeProperty = edgeConf.createDefaultCustomEdgeProperty(edgeManagerDescriptor);

  return DAG.create("CrossProduct").addVertex(v1).addVertex(v2).addVertex(v3)
    .addEdge(Edge.create(v1, v3, edgeProperty)).addEdge(Edge.create(v2, v3, edgeProperty));
}
 

public static InputSpec getInputSpecFromProto(IOSpecProto inputSpecProto) {
  InputDescriptor inputDescriptor = null;
  if (inputSpecProto.hasIoDescriptor()) {
    inputDescriptor =
        DagTypeConverters.convertInputDescriptorFromDAGPlan(inputSpecProto.getIoDescriptor());
  }
  InputSpec inputSpec = new InputSpec(inputSpecProto.getConnectedVertexName(), inputDescriptor,
      inputSpecProto.getPhysicalEdgeCount());
  return inputSpec;
}
 

/**
 * This is a convenience method for the typical usage of this edge, and creates an instance of
 * {@link org.apache.tez.dag.api.EdgeProperty} which is likely to be used. </p>
 * If custom edge properties are required, the methods to get the relevant payloads should be
 * used. </p>
 * * In this case - DataMovementType.SCATTER_GATHER, EdgeProperty.DataSourceType.PERSISTED,
 * EdgeProperty.SchedulingType.SEQUENTIAL
 *
 * @return an {@link org.apache.tez.dag.api.EdgeProperty} instance
 */
public EdgeProperty createDefaultEdgeProperty() {
  EdgeProperty edgeProperty = EdgeProperty.create(EdgeProperty.DataMovementType.SCATTER_GATHER,
      EdgeProperty.DataSourceType.PERSISTED, EdgeProperty.SchedulingType.SEQUENTIAL,
      OutputDescriptor.create(
          getOutputClassName()).setUserPayload(getOutputPayload()),
      InputDescriptor.create(
          getInputClassName()).setUserPayload(getInputPayload()));
  Utils.setEdgePropertyHistoryText(this, edgeProperty);
  return edgeProperty;
}
 

public void runInputInitializers(List<RootInputLeafOutput<InputDescriptor, InputInitializerDescriptor>>
    inputs) throws TezException {
  for (RootInputLeafOutput<InputDescriptor, InputInitializerDescriptor> input : inputs) {

    InputInitializerContext context =
        new TezRootInputInitializerContextImpl(input, vertex, appContext, this);

    InputInitializer initializer;
    try {
      TezUtilsInternal.setHadoopCallerContext(appContext.getHadoopShim(), vertex.getVertexId());
      initializer = createInitializer(input, context);
    } finally {
      appContext.getHadoopShim().clearHadoopCallerContext();
    }

    InitializerWrapper initializerWrapper =
        new InitializerWrapper(input, initializer, context, vertex, entityStateTracker, appContext);

    // Register pending vertex update registrations
    List<VertexUpdateRegistrationHolder> vertexUpdateRegistrations = pendingVertexRegistrations.removeAll(input.getName());
    if (vertexUpdateRegistrations != null) {
      for (VertexUpdateRegistrationHolder h : vertexUpdateRegistrations) {
        initializerWrapper.registerForVertexStateUpdates(h.vertexName, h.stateSet);
      }
    }

    initializerMap.put(input.getName(), initializerWrapper);
    ListenableFuture<List<Event>> future = executor
        .submit(new InputInitializerCallable(initializerWrapper, dagUgi, appContext));
    Futures.addCallback(future, createInputInitializerCallback(initializerWrapper), GuavaShim.directExecutor());
  }
}
 

@Override
public void onRootVertexInitialized(String inputName,
    InputDescriptor inputDescriptor, List<Event> events) {
  List<InputDataInformationEvent> inputInfoEvents = new ArrayList<InputDataInformationEvent>();
  for (Event event: events) {
    if (event instanceof InputDataInformationEvent) {
      inputInfoEvents.add((InputDataInformationEvent)event);
    }
  }
  getContext().addRootInputEvents(inputName, inputInfoEvents);
}
 

public static InputDescriptor getInputDesc(UserPayload payload) {
  InputDescriptor desc = InputDescriptor.create(TestInput.class.getName());
  if (payload != null) {
    desc.setUserPayload(payload);
  }
  return desc;
}
 

@Test (timeout=5000)
public void testBasicScatterGather() throws Exception {
  HashMap<String, EdgeProperty> mockInputVertices = 
      new HashMap<String, EdgeProperty>();
  String mockSrcVertexId1 = "Vertex1";
  EdgeProperty eProp1 = EdgeProperty.create(
      EdgeProperty.DataMovementType.SCATTER_GATHER,
      EdgeProperty.DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create("out"),
      InputDescriptor.create("in"));
  
  String mockManagedVertexId = "Vertex";
  
  VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
  when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
  when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
  when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(2);
  when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(3);
  mockInputVertices.put(mockSrcVertexId1, eProp1);

  InputReadyVertexManager manager = new InputReadyVertexManager(mockContext);
  manager.initialize();
  verify(mockContext, times(1)).vertexReconfigurationPlanned();
  // source vertex configured
  manager.onVertexStateUpdated(new VertexStateUpdate(mockSrcVertexId1, VertexState.CONFIGURED));
  verify(mockContext, times(1)).doneReconfiguringVertex();
  verify(mockContext, times(0)).scheduleTasks(requestCaptor.capture());
  // then own vertex started
  manager.onVertexStarted(Collections.singletonList(
      TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 0)));
  manager.onSourceTaskCompleted(
      TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 1));
  verify(mockContext, times(0)).scheduleTasks(anyList());
  manager.onSourceTaskCompleted(
      TestShuffleVertexManager.createTaskAttemptIdentifier(mockSrcVertexId1, 2));
  verify(mockContext, times(1)).scheduleTasks(requestCaptor.capture());
  Assert.assertEquals(2, requestCaptor.getValue().size());
}
 

InitializerWrapper(RootInputLeafOutput<InputDescriptor, InputInitializerDescriptor> input,
                   InputInitializer initializer, InputInitializerContext context,
                   Vertex vertex, StateChangeNotifier stateChangeNotifier,
                   AppContext appContext) {
  this.input = input;
  this.initializer = initializer;
  this.context = context;
  this.vertexLogIdentifier = vertex.getLogIdentifier();
  this.vertexId = vertex.getVertexId();
  this.stateChangeNotifier = stateChangeNotifier;
  this.appContext = appContext;
}
 

@Ignore
@Test (timeout = 600000)
public void testMemoryBroadcast() throws Exception {
  DAG dag = DAG.create("testMemoryBroadcast");
  Vertex vA = Vertex.create("A", ProcessorDescriptor.create("Proc.class"), numTasks);
  Vertex vB = Vertex.create("B", ProcessorDescriptor.create("Proc.class"), numTasks);
  dag.addVertex(vA)
      .addVertex(vB)
      .addEdge(
          Edge.create(vA, vB, EdgeProperty.create(DataMovementType.BROADCAST,
              DataSourceType.PERSISTED, SchedulingType.SEQUENTIAL,
              OutputDescriptor.create("Out"), InputDescriptor.create("In"))));
  testMemory(dag, true);
}
 

/**
 * This is a convenience method for creating an Edge descriptor based on the specified
 * EdgeManagerDescriptor.
 *
 * @param edgeManagerDescriptor the custom edge specification
 * @return an {@link org.apache.tez.dag.api.EdgeProperty} instance
 */
public EdgeProperty createDefaultCustomEdgeProperty(EdgeManagerPluginDescriptor edgeManagerDescriptor) {
  Objects.requireNonNull(edgeManagerDescriptor, "EdgeManagerDescriptor cannot be null");
  EdgeProperty edgeProperty =
      EdgeProperty.create(edgeManagerDescriptor, EdgeProperty.DataSourceType.PERSISTED,
          EdgeProperty.SchedulingType.SEQUENTIAL,
          OutputDescriptor.create(getOutputClassName()).setUserPayload(getOutputPayload()),
          InputDescriptor.create(getInputClassName()).setUserPayload(getInputPayload()));
  Utils.setEdgePropertyHistoryText(this, edgeProperty);
  return edgeProperty;
}
 

@Override
public void readFields(DataInput in) throws IOException {
  sourceVertexName = StringInterner.weakIntern(in.readUTF());
  physicalEdgeCount = in.readInt();
  inputDescriptor = new InputDescriptor();
  inputDescriptor.readFields(in);
}
 

@Private
public TezInputContextImpl(Configuration conf, String[] workDirs,
                           int appAttemptNumber,
                           TezUmbilical tezUmbilical, String dagName, 
                           String taskVertexName, String sourceVertexName,
                           int vertexParallelism, TezTaskAttemptID taskAttemptID,
                           int inputIndex, @Nullable UserPayload userPayload,
                           LogicalIOProcessorRuntimeTask runtimeTask,
                           Map<String, ByteBuffer> serviceConsumerMetadata,
                           Map<String, String> auxServiceEnv, MemoryDistributor memDist,
                           InputDescriptor inputDescriptor, Map<String, LogicalInput> inputs,
                           InputReadyTracker inputReadyTracker, ObjectRegistry objectRegistry,
                           ExecutionContext ExecutionContext, long memAvailable,
                           TezExecutors sharedExecutor) {
  super(conf, workDirs, appAttemptNumber, dagName, taskVertexName,
      vertexParallelism, taskAttemptID, wrapCounters(runtimeTask,
      taskVertexName, sourceVertexName, conf), runtimeTask, tezUmbilical,
      serviceConsumerMetadata, auxServiceEnv, memDist, inputDescriptor,
      objectRegistry, ExecutionContext, memAvailable, sharedExecutor);
  Objects.requireNonNull(inputIndex, "inputIndex is null");
  Objects.requireNonNull(sourceVertexName, "sourceVertexName is null");
  Objects.requireNonNull(inputs, "input map is null");
  Objects.requireNonNull(inputReadyTracker, "inputReadyTracker is null");
  this.userPayload = userPayload;
  this.inputIndex = inputIndex;
  this.sourceVertexName = sourceVertexName;
  this.sourceInfo = new EventMetaData(
      EventProducerConsumerType.INPUT, taskVertexName, sourceVertexName,
      taskAttemptID);
  this.inputs = inputs;
  this.inputReadyTracker = inputReadyTracker;
  runtimeTask.getTaskStatistics().addIO(sourceVertexName);
  statsReporter = new InputStatisticsReporterImpl();
}
 

@Test(timeout = 5000)
public void testScalingNoProcessor() throws TezException {
  MemoryDistributor dist = new MemoryDistributor(2, 1, conf);
  
  dist.setJvmMemory(10000l);

  // First request
  MemoryUpdateCallbackForTest e1Callback = new MemoryUpdateCallbackForTest();
  InputContext e1InputContext1 = createTestInputContext();
  InputDescriptor e1InDesc1 = createTestInputDescriptor();
  dist.requestMemory(10000, e1Callback, e1InputContext1, e1InDesc1);
  
  // Second request
  MemoryUpdateCallbackForTest e2Callback = new MemoryUpdateCallbackForTest();
  InputContext e2InputContext2 = createTestInputContext();
  InputDescriptor e2InDesc2 = createTestInputDescriptor();
  dist.requestMemory(10000, e2Callback, e2InputContext2, e2InDesc2);
  
  // Third request - output
  MemoryUpdateCallbackForTest e3Callback = new MemoryUpdateCallbackForTest();
  OutputContext e3OutputContext1 = createTestOutputContext();
  OutputDescriptor e3OutDesc2 = createTestOutputDescriptor();
  dist.requestMemory(5000, e3Callback, e3OutputContext1, e3OutDesc2);
  
  dist.makeInitialAllocations();
  
  // Total available: 70% of 10K = 7000
  // 3 requests - 10K, 10K, 5K
  // Scale down to - 2800, 2800, 1400
  assertEquals(2800, e1Callback.assigned);
  assertEquals(2800, e2Callback.assigned);
  assertEquals(1400, e3Callback.assigned);
}
 

public void onRootVertexInitialized(String inputName,
    InputDescriptor inputDescriptor, List<Event> events) throws AMUserCodeException {
  if (LOG.isDebugEnabled()) {
    LOG.debug("vertex:" + managedVertex.getLogIdentifier() + "; enqueueing onRootVertexInitialized"
        + " on input:" + inputName + ", current task events size is " + rootInputInitEventQueue.size());
  }
  enqueueAndScheduleNextEvent(new VertexManagerEventRootInputInitialized(inputName,
      inputDescriptor, events));
}
 

private LogicalInput createInputFromDescriptor(InputDescriptor inputDesc) {
  Input input = ReflectionUtils.createClazzInstance(inputDesc.getClassName());
  if (!(input instanceof LogicalInput)) {
    throw new TezUncheckedException(inputDesc.getClass().getName()
        + " is not a sub-type of LogicalInput."
        + " Only LogicalInput sub-types supported by LogicalIOProcessor.");
  }
  return (LogicalInput)input;
}
 

/**
 * only addRootInputEvents when it is "input2", otherwise just cache it.
 */
@Override
public void onRootVertexInitialized(String inputName,
    InputDescriptor inputDescriptor, List<Event> events) {
  cachedEventMap.put(inputName, events);
  if (inputName.equals("input2")) {
    for (Map.Entry<String, List<Event>> entry : cachedEventMap.entrySet()) {
      List<InputDataInformationEvent> riEvents = Lists.newLinkedList();
      for (Event event : events) {
        riEvents.add((InputDataInformationEvent)event);
      }
      getContext().addRootInputEvents(entry.getKey(), riEvents);
    }
  }
}
 

private DataSourceDescriptor createDistributorDataSource() throws IOException {
  InputSplitInfo inputSplitInfo;
  setupBasicConf(conf);
  try {
    inputSplitInfo = MRInputHelpers.generateInputSplitsToMem(conf, false, true, 0);
  } catch (Exception e) {
    throw new TezUncheckedException(e);
  }
  MRHelpers.translateMRConfToTez(conf);

  UserPayload payload = MRInputHelpersInternal.createMRInputPayload(conf,
      inputSplitInfo.getSplitsProto());
  Credentials credentials = null;
  if (getCredentialsForSourceFilesystem && inputSplitInfo.getCredentials() != null) {
    credentials = inputSplitInfo.getCredentials();
  }
  DataSourceDescriptor ds = DataSourceDescriptor.create(
      InputDescriptor.create(inputClassName).setUserPayload(payload),
      InputInitializerDescriptor.create(MRInputSplitDistributor.class.getName()),
      inputSplitInfo.getNumTasks(), credentials,
      VertexLocationHint.create(inputSplitInfo.getTaskLocationHints()), null);
  if (conf.getBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_CONVERT_USER_PAYLOAD_TO_HISTORY_TEXT,
      TezRuntimeConfiguration.TEZ_RUNTIME_CONVERT_USER_PAYLOAD_TO_HISTORY_TEXT_DEFAULT)) {
    ds.getInputDescriptor().setHistoryText(TezUtils.convertToHistoryText(conf));
  }

  return ds;
}
 

@Test(timeout = 5000)
public void testEventsFromMultipleInputs() throws IOException {

  VertexManagerPluginContext context = mock(VertexManagerPluginContext.class);
  TezConfiguration conf = new TezConfiguration();
  UserPayload vertexPayload = TezUtils.createUserPayloadFromConf(conf);
  doReturn("vertex1").when(context).getVertexName();
  doReturn(1).when(context).getVertexNumTasks(eq("vertex1"));
  doReturn(vertexPayload).when(context).getUserPayload();

  RootInputVertexManager rootInputVertexManager = new RootInputVertexManager(context);
  rootInputVertexManager.initialize();

  InputDescriptor id1 = mock(InputDescriptor.class);
  List<Event> events1 = new LinkedList<Event>();
  InputDataInformationEvent diEvent11 = InputDataInformationEvent.createWithSerializedPayload(0,
      null);
  events1.add(diEvent11);
  rootInputVertexManager.onRootVertexInitialized("input1", id1, events1);
  // All good so far, single input only.

  InputDescriptor id2 = mock(InputDescriptor.class);
  List<Event> events2 = new LinkedList<Event>();
  InputDataInformationEvent diEvent21 = InputDataInformationEvent.createWithSerializedPayload(0,
      null);
  events2.add(diEvent21);
  try {
    // Should fail due to second input
    rootInputVertexManager.onRootVertexInitialized("input2", id2, events2);
    fail("Expecting failure in case of multiple inputs attempting to send events");
  } catch (IllegalStateException e) {
    assertTrue(e.getMessage().startsWith(
        "RootInputVertexManager cannot configure multiple inputs. Use a custom VertexManager"));
  }
}
 

@Override
public void onRootVertexInitialized(String inputName, InputDescriptor inputDescriptor,
    List<Event> events) {
  Map<String, RootInputSpecUpdate> map = new HashMap<String, RootInputSpecUpdate>();
  if (context.getUserPayload()[0] == 0) {
    map.put("input3", RootInputSpecUpdate.createAllTaskRootInputSpecUpdate(4));
  } else {
    List<Integer> pInputList = new LinkedList<Integer>();
    for (int i = 1; i <= NUM_TASKS; i++) {
      pInputList.add(i);
    }
    map.put("input4", RootInputSpecUpdate.createPerTaskRootInputSpecUpdate(pInputList));
  }
  context.setVertexParallelism(NUM_TASKS, null, null, map);
}
 

@Test
public void testEventsFromMultipleInputs() {

  VertexManagerPluginContext context = mock(VertexManagerPluginContext.class);
  doReturn("vertex1").when(context).getVertexName();
  doReturn(1).when(context).getVertexNumTasks(eq("vertex1"));

  RootInputVertexManager rootInputVertexManager = new RootInputVertexManager();
  rootInputVertexManager.initialize(context);

  InputDescriptor id1 = mock(InputDescriptor.class);
  List<Event> events1 = new LinkedList<Event>();
  RootInputDataInformationEvent diEvent11 = new RootInputDataInformationEvent(0, null);
  events1.add(diEvent11);
  rootInputVertexManager.onRootVertexInitialized("input1", id1, events1);
  // All good so far, single input only.

  InputDescriptor id2 = mock(InputDescriptor.class);
  List<Event> events2 = new LinkedList<Event>();
  RootInputDataInformationEvent diEvent21 = new RootInputDataInformationEvent(0, null);
  events2.add(diEvent21);
  try {
    // Should fail due to second input
    rootInputVertexManager.onRootVertexInitialized("input2", id2, events2);
    fail("Expecting failure in case of multiple inputs attempting to send events");
  } catch (IllegalStateException e) {
    assertTrue(e.getMessage().startsWith(
        "RootInputVertexManager cannot configure multiple inputs. Use a custom VertexManager"));
  }
}
 

/**
 * TEZ-1647
 * custom vertex manager generates events only when both i1 and i2 are initialized.
 * @throws Exception
 */
@Test(timeout = 5000)
public void testOnRootVertexInitialized2() throws Exception {
  VertexManager vm =
      new VertexManager(
          VertexManagerPluginDescriptor.create(CustomVertexManager.class
              .getName()), UserGroupInformation.getCurrentUser(),
              mockVertex, mockAppContext, mock(StateChangeNotifier.class));
  vm.initialize();
  InputDescriptor id1 = mock(InputDescriptor.class);
  List<Event> events1 = new LinkedList<Event>();
  InputDataInformationEvent diEvent1 =
      InputDataInformationEvent.createWithSerializedPayload(0, null);
  events1.add(diEvent1);

  // do not call context.addRootInputEvents, just cache the TezEvent
  vm.onRootVertexInitialized("input1", id1, events1);
  verify(mockHandler, times(1)).handle(requestCaptor.capture());
  List<TezEvent> tezEventsAfterInput1 = requestCaptor.getValue().getEvents();
  assertEquals(0, tezEventsAfterInput1.size());
  
  InputDescriptor id2 = mock(InputDescriptor.class);
  List<Event> events2 = new LinkedList<Event>();
  InputDataInformationEvent diEvent2 =
      InputDataInformationEvent.createWithSerializedPayload(0, null);
  events2.add(diEvent2);
  // call context.addRootInputEvents(input1), context.addRootInputEvents(input2)
  vm.onRootVertexInitialized("input2", id2, events2);
  verify(mockHandler, times(2)).handle(requestCaptor.capture());
  List<TezEvent> tezEventsAfterInput2 = requestCaptor.getValue().getEvents();
  assertEquals(2, tezEventsAfterInput2.size());

  // also verify the EventMetaData
  Set<String> edgeVertexSet = new HashSet<String>();
  for (TezEvent tezEvent : tezEventsAfterInput2) {
    edgeVertexSet.add(tezEvent.getDestinationInfo().getEdgeVertexName());
  }
  assertEquals(Sets.newHashSet("input1","input2"), edgeVertexSet);
}
 

@Test
public void testZeroTasksSendsConfigured() throws IOException {
  Configuration conf = new Configuration();
  ShuffleVertexManagerBase manager = null;

  HashMap<String, EdgeProperty> mockInputVertices = new HashMap<String, EdgeProperty>();
  String r1 = "R1";
  EdgeProperty eProp1 = EdgeProperty.create(
      EdgeProperty.DataMovementType.SCATTER_GATHER,
      EdgeProperty.DataSourceType.PERSISTED,
      SchedulingType.SEQUENTIAL,
      OutputDescriptor.create("out"),
      InputDescriptor.create("in"));

  final String mockManagedVertexId = "R2";
  mockInputVertices.put(r1, eProp1);

  final VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
  when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
  when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
  when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(0);

  // check initialization
  manager = createManager(conf, mockContext, 0.001f, 0.001f);

  final List<Integer> scheduledTasks = Lists.newLinkedList();
  doAnswer(new ScheduledTasksAnswer(scheduledTasks)).when(
      mockContext).scheduleTasks(anyList());

  manager.onVertexStarted(emptyCompletions);
  manager.onVertexStateUpdated(new VertexStateUpdate(r1, VertexState.CONFIGURED));
  Assert.assertEquals(1, manager.bipartiteSources);
  Assert.assertEquals(0, manager.numBipartiteSourceTasksCompleted);
  Assert.assertEquals(0, manager.totalNumBipartiteSourceTasks);
  Assert.assertEquals(0, manager.pendingTasks.size()); // no tasks scheduled
  Assert.assertEquals(0, scheduledTasks.size());
  verify(mockContext).doneReconfiguringVertex();
}