/**
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.tez.test.dag;
import com.google.common.collect.Lists;
import com.google.common.primitives.Ints;
import java.nio.ByteBuffer;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.tez.common.TezUtils;
import org.apache.tez.dag.api.DAG;
import org.apache.tez.dag.api.Edge;
import org.apache.tez.dag.api.EdgeProperty;
import org.apache.tez.dag.api.EdgeProperty.DataMovementType;
import org.apache.tez.dag.api.EdgeProperty.DataSourceType;
import org.apache.tez.dag.api.EdgeProperty.SchedulingType;
import org.apache.tez.dag.api.InputDescriptor;
import org.apache.tez.dag.api.UserPayload;
import org.apache.tez.dag.api.Vertex;
import org.apache.tez.dag.api.VertexManagerPlugin;
import org.apache.tez.dag.api.VertexManagerPluginContext;
import org.apache.tez.dag.api.VertexManagerPluginDescriptor;
import org.apache.tez.dag.api.VertexManagerPluginContext.TaskWithLocationHint;
import org.apache.tez.dag.api.client.VertexStatus.State;
import org.apache.tez.runtime.api.AbstractLogicalInput;
import org.apache.tez.runtime.api.AbstractLogicalOutput;
import org.apache.tez.runtime.api.Event;
import org.apache.tez.runtime.api.MemoryUpdateCallback;
import org.apache.tez.runtime.api.OutputCommitter;
import org.apache.tez.runtime.api.OutputCommitterContext;
import org.apache.tez.runtime.api.Reader;
import org.apache.tez.runtime.api.InputContext;
import org.apache.tez.runtime.api.OutputContext;
import org.apache.tez.runtime.api.InputInitializer;
import org.apache.tez.runtime.api.InputInitializerContext;
import org.apache.tez.runtime.api.Writer;
import org.apache.tez.runtime.api.events.InputInitializerEvent;
import org.apache.tez.runtime.api.events.VertexManagerEvent;
import org.apache.tez.test.TestInput;
import org.apache.tez.test.TestOutput;
import org.apache.tez.test.TestProcessor;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.atomic.AtomicInteger;
public class MultiAttemptDAG {
private static final Log LOG =
LogFactory.getLog(MultiAttemptDAG.class);
static Resource defaultResource = Resource.newInstance(100, 0);
public static String MULTI_ATTEMPT_DAG_VERTEX_NUM_TASKS =
"tez.multi-attempt-dag.vertex.num-tasks";
public static int MULTI_ATTEMPT_DAG_VERTEX_NUM_TASKS_DEFAULT = 2;
public static String MULTI_ATTEMPT_DAG_USE_FAILING_COMMITTER =
"tez.multi-attempt-dag.use-failing-committer";
public static boolean MULTI_ATTEMPT_DAG_USE_FAILING_COMMITTER_DEFAULT = false;
public static class FailOnAttemptVertexManagerPlugin extends VertexManagerPlugin {
private int numSourceTasks = 0;
private AtomicInteger numCompletions = new AtomicInteger();
private boolean tasksScheduled = false;
public FailOnAttemptVertexManagerPlugin(VertexManagerPluginContext context) {
super(context);
}
@Override
public void initialize() {
for (String input :
getContext().getInputVertexEdgeProperties().keySet()) {
LOG.info("Adding sourceTasks for Vertex " + input);
numSourceTasks += getContext().getVertexNumTasks(input);
LOG.info("Current numSourceTasks=" + numSourceTasks);
}
}
@Override
public void onVertexStarted(Map<String, List<Integer>> completions) {
if (completions != null) {
for (Entry<String, List<Integer>> entry : completions.entrySet()) {
LOG.info("Received completion events on vertexStarted"
+ ", vertex=" + entry.getKey()
+ ", completions=" + entry.getValue().size());
numCompletions.addAndGet(entry.getValue().size());
}
}
maybeScheduleTasks();
}
private synchronized void maybeScheduleTasks() {
if (numCompletions.get() >= numSourceTasks
&& !tasksScheduled) {
tasksScheduled = true;
String payload = new String(getContext().getUserPayload().deepCopyAsArray());
int successAttemptId = Integer.valueOf(payload);
LOG.info("Checking whether to crash AM or schedule tasks"
+ ", successfulAttemptID=" + successAttemptId
+ ", currentAttempt=" + getContext().getDAGAttemptNumber());
if (successAttemptId > getContext().getDAGAttemptNumber()) {
Runtime.getRuntime().halt(-1);
} else if (successAttemptId == getContext().getDAGAttemptNumber()) {
LOG.info("Scheduling tasks for vertex=" + getContext().getVertexName());
int numTasks = getContext().getVertexNumTasks(getContext().getVertexName());
List<TaskWithLocationHint> scheduledTasks = Lists.newArrayListWithCapacity(numTasks);
for (int i=0; i<numTasks; ++i) {
scheduledTasks.add(new TaskWithLocationHint(new Integer(i), null));
}
getContext().scheduleVertexTasks(scheduledTasks);
}
}
}
@Override
public void onSourceTaskCompleted(String srcVertexName, Integer taskId) {
LOG.info("Received completion events for source task"
+ ", vertex=" + srcVertexName
+ ", taskIdx=" + taskId);
numCompletions.incrementAndGet();
maybeScheduleTasks();
}
@Override
public void onVertexManagerEventReceived(VertexManagerEvent vmEvent) {
// Nothing to do
}
@Override
public void onRootVertexInitialized(String inputName, InputDescriptor inputDescriptor, List<Event> events) {
// Do nothing
}
}
public static class FailingOutputCommitter extends OutputCommitter {
boolean failOnCommit = false;
public FailingOutputCommitter(OutputCommitterContext committerContext) {
super(committerContext);
}
@Override
public void initialize() throws Exception {
FailingOutputCommitterConfig config = new
FailingOutputCommitterConfig();
config.fromUserPayload(
getContext().getOutputUserPayload().deepCopyAsArray());
failOnCommit = config.failOnCommit;
}
@Override
public void setupOutput() throws Exception {
}
@Override
public void commitOutput() throws Exception {
if (failOnCommit) {
LOG.info("Committer causing AM to shutdown");
Runtime.getRuntime().halt(-1);
}
}
@Override
public void abortOutput(State finalState) throws Exception {
}
public static class FailingOutputCommitterConfig {
boolean failOnCommit;
public FailingOutputCommitterConfig() {
this(false);
}
public FailingOutputCommitterConfig(boolean failOnCommit) {
this.failOnCommit = failOnCommit;
}
public byte[] toUserPayload() {
return Ints.toByteArray((failOnCommit ? 1 : 0));
}
public void fromUserPayload(byte[] userPayload) {
int failInt = Ints.fromByteArray(userPayload);
if (failInt == 0) {
failOnCommit = false;
} else {
failOnCommit = true;
}
}
}
}
public static class FailingInputInitializer extends InputInitializer {
public FailingInputInitializer(InputInitializerContext initializerContext) {
super(initializerContext);
}
@Override
public List<Event> initialize() throws Exception {
try {
Thread.sleep(2000l);
} catch (InterruptedException e) {
// Ignore
}
if (getContext().getDAGAttemptNumber() == 1) {
LOG.info("Shutting down the AM in 1st attempt");
Runtime.getRuntime().halt(-1);
}
return null;
}
@Override
public void handleInputInitializerEvent(List<InputInitializerEvent> events) throws
Exception {
throw new UnsupportedOperationException("Not supported");
}
}
public static class NoOpInput extends AbstractLogicalInput {
public NoOpInput(InputContext inputContext, int numPhysicalInputs) {
super(inputContext, numPhysicalInputs);
}
@Override
public List<Event> initialize() throws Exception {
getContext().requestInitialMemory(1l, new MemoryUpdateCallback() {
@Override
public void memoryAssigned(long assignedSize) {}
});
return null;
}
@Override
public void start() throws Exception {
}
@Override
public Reader getReader() throws Exception {
return null;
}
@Override
public void handleEvents(List<Event> inputEvents) throws Exception {
}
@Override
public List<Event> close() throws Exception {
return null;
}
}
public static class NoOpOutput extends AbstractLogicalOutput {
public NoOpOutput(OutputContext outputContext,
int numPhysicalOutputs) {
super(outputContext, numPhysicalOutputs);
}
@Override
public List<Event> initialize() throws Exception {
getContext().requestInitialMemory(1l, new MemoryUpdateCallback() {
@Override
public void memoryAssigned(long assignedSize) {
}
});
return null;
}
@Override
public void start() throws Exception {
}
@Override
public Writer getWriter() throws Exception {
return null;
}
@Override
public void handleEvents(List<Event> outputEvents) {
}
@Override
public List<Event> close() throws Exception {
return null;
}
}
public static DAG createDAG(String name,
Configuration conf) throws Exception {
UserPayload payload = UserPayload.create(null);
int taskCount = MULTI_ATTEMPT_DAG_VERTEX_NUM_TASKS_DEFAULT;
if (conf != null) {
taskCount = conf.getInt(MULTI_ATTEMPT_DAG_VERTEX_NUM_TASKS, MULTI_ATTEMPT_DAG_VERTEX_NUM_TASKS_DEFAULT);
payload = TezUtils.createUserPayloadFromConf(conf);
}
DAG dag = DAG.create(name);
Vertex v1 = Vertex.create("v1", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
Vertex v2 = Vertex.create("v2", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
Vertex v3 = Vertex.create("v3", TestProcessor.getProcDesc(payload), taskCount, defaultResource);
// Make each vertex manager fail on appropriate attempt
v1.setVertexManagerPlugin(VertexManagerPluginDescriptor.create(
FailOnAttemptVertexManagerPlugin.class.getName())
.setUserPayload(UserPayload.create(ByteBuffer.wrap(new String("1").getBytes()))));
v2.setVertexManagerPlugin(VertexManagerPluginDescriptor.create(
FailOnAttemptVertexManagerPlugin.class.getName())
.setUserPayload(UserPayload.create(ByteBuffer.wrap(new String("2").getBytes()))));
v3.setVertexManagerPlugin(VertexManagerPluginDescriptor.create(
FailOnAttemptVertexManagerPlugin.class.getName())
.setUserPayload(UserPayload.create(ByteBuffer.wrap(new String("3").getBytes()))));
dag.addVertex(v1).addVertex(v2).addVertex(v3);
dag.addEdge(Edge.create(v1, v2,
EdgeProperty.create(DataMovementType.SCATTER_GATHER,
DataSourceType.PERSISTED,
SchedulingType.SEQUENTIAL,
TestOutput.getOutputDesc(payload),
TestInput.getInputDesc(payload))));
dag.addEdge(Edge.create(v2, v3,
EdgeProperty.create(DataMovementType.SCATTER_GATHER,
DataSourceType.PERSISTED,
SchedulingType.SEQUENTIAL,
TestOutput.getOutputDesc(payload),
TestInput.getInputDesc(payload))));
return dag;
}
public static DAG createDAG(Configuration conf) throws Exception {
return createDAG("SimpleVTestDAG", conf);
}
}