You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tez.apache.org by bi...@apache.org on 2014/04/22 19:48:44 UTC
git commit: TEZ-1018. VertexManagerPluginContext should enable
assigning locality to scheduled tasks (bikas)
Repository: incubator-tez
Updated Branches:
refs/heads/master 617c8b8ea -> 0a681b34e
TEZ-1018. VertexManagerPluginContext should enable assigning locality to scheduled tasks (bikas)
Project: http://git-wip-us.apache.org/repos/asf/incubator-tez/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-tez/commit/0a681b34
Tree: http://git-wip-us.apache.org/repos/asf/incubator-tez/tree/0a681b34
Diff: http://git-wip-us.apache.org/repos/asf/incubator-tez/diff/0a681b34
Branch: refs/heads/master
Commit: 0a681b34e154bf622c68f753820153d7ee77f4c4
Parents: 617c8b8
Author: Bikas Saha <bi...@apache.org>
Authored: Tue Apr 22 10:48:38 2014 -0700
Committer: Bikas Saha <bi...@apache.org>
Committed: Tue Apr 22 10:48:38 2014 -0700
----------------------------------------------------------------------
CHANGES.txt | 2 +
.../main/java/org/apache/tez/dag/api/DAG.java | 4 +
.../apache/tez/dag/api/VertexLocationHint.java | 37 ++-
.../tez/dag/api/VertexManagerPluginContext.java | 35 ++-
.../dag/api/VertexManagerPluginDescriptor.java | 4 +-
.../java/org/apache/tez/dag/app/dag/Task.java | 1 +
.../org/apache/tez/dag/app/dag/TaskAttempt.java | 6 +
.../java/org/apache/tez/dag/app/dag/Vertex.java | 3 +-
.../dag/impl/ImmediateStartVertexManager.java | 8 +-
.../app/dag/impl/RootInputVertexManager.java | 6 +-
.../tez/dag/app/dag/impl/TaskAttemptImpl.java | 14 +-
.../apache/tez/dag/app/dag/impl/TaskImpl.java | 5 +-
.../apache/tez/dag/app/dag/impl/VertexImpl.java | 66 +++--
.../tez/dag/app/dag/impl/VertexManager.java | 18 +-
.../dag/app/rm/TaskSchedulerEventHandler.java | 23 +-
.../tez/dag/app/dag/impl/TestVertexImpl.java | 18 +-
.../examples/BroadcastAndOneToOneExample.java | 287 +++++++++++++++++++
.../vertexmanager/InputReadyVertexManager.java | 45 ++-
.../vertexmanager/ShuffleVertexManager.java | 10 +-
.../TestInputReadyVertexManager.java | 54 +++-
.../vertexmanager/TestShuffleVertexManager.java | 11 +-
.../apache/tez/mapreduce/TestMRRJobsDAGApi.java | 16 +-
.../apache/tez/test/dag/MultiAttemptDAG.java | 7 +-
23 files changed, 592 insertions(+), 88 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/CHANGES.txt
----------------------------------------------------------------------
diff --git a/CHANGES.txt b/CHANGES.txt
index ac79072..fefb26b 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -7,6 +7,8 @@ INCOMPATIBLE CHANGES
TEZ-960. VertexManagerPluginContext::getTotalAVailableResource() changed to
VertexManagerPluginContext::getTotalAvailableResource()
TEZ-1025. Rename tez.am.max.task.attempts to tez.am.task.max.failed.attempts
+ TEZ-1018. VertexManagerPluginContext should enable assigning locality to
+ scheduled tasks
Release 0.4.0-incubating: 2014-04-05
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-api/src/main/java/org/apache/tez/dag/api/DAG.java
----------------------------------------------------------------------
diff --git a/tez-api/src/main/java/org/apache/tez/dag/api/DAG.java b/tez-api/src/main/java/org/apache/tez/dag/api/DAG.java
index fb82f08..5e9c375 100644
--- a/tez-api/src/main/java/org/apache/tez/dag/api/DAG.java
+++ b/tez-api/src/main/java/org/apache/tez/dag/api/DAG.java
@@ -547,6 +547,10 @@ public class DAG { // FIXME rename to Topology
for (TaskLocationHint hint : vertex.getTaskLocationsHint().getTaskLocationHints()) {
PlanTaskLocationHint.Builder taskLocationHintBuilder = PlanTaskLocationHint.newBuilder();
+ if (hint.getAffinitizedContainer() != null) {
+ throw new TezUncheckedException(
+ "Container affinity may not be specified via the DAG API");
+ }
if (hint.getDataLocalHosts() != null) {
taskLocationHintBuilder.addAllHost(hint.getDataLocalHosts());
}
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-api/src/main/java/org/apache/tez/dag/api/VertexLocationHint.java
----------------------------------------------------------------------
diff --git a/tez-api/src/main/java/org/apache/tez/dag/api/VertexLocationHint.java b/tez-api/src/main/java/org/apache/tez/dag/api/VertexLocationHint.java
index 8c0fe35..979d363 100644
--- a/tez-api/src/main/java/org/apache/tez/dag/api/VertexLocationHint.java
+++ b/tez-api/src/main/java/org/apache/tez/dag/api/VertexLocationHint.java
@@ -22,6 +22,10 @@ import java.util.Collections;
import java.util.List;
import java.util.Set;
+import org.apache.hadoop.yarn.api.records.ContainerId;
+
+import com.google.common.base.Preconditions;
+
public class VertexLocationHint {
private final List<TaskLocationHint> taskLocationHints;
@@ -73,11 +77,26 @@ public class VertexLocationHint {
public static class TaskLocationHint {
+ // Container to which to affinitize
+ ContainerId containerId;
// Host names if any to be used
- private final Set<String> hosts;
+ private Set<String> hosts;
// Rack names if any to be used
- private final Set<String> racks;
+ private Set<String> racks;
+ /**
+ * Provide a container to which the task may be affinitized
+ * Tez will try to run the task on that container or node local to it
+ */
+ public TaskLocationHint(ContainerId containerId) {
+ Preconditions.checkNotNull(containerId);
+ this.containerId = containerId;
+ }
+
+ /**
+ * Provides nodes and racks at which the task may be executed.
+ * Tez will try to execute the task at those locations
+ */
public TaskLocationHint(Set<String> hosts, Set<String> racks) {
if (hosts != null) {
this.hosts = Collections.unmodifiableSet(hosts);
@@ -91,6 +110,10 @@ public class VertexLocationHint {
}
}
+ public ContainerId getAffinitizedContainer() {
+ return containerId;
+ }
+
public Set<String> getDataLocalHosts() {
return hosts;
}
@@ -109,6 +132,9 @@ public class VertexLocationHint {
result = ( racks != null) ?
prime * result + racks.hashCode() :
result + prime;
+ result = ( containerId != null) ?
+ prime * result + containerId.hashCode() :
+ result + prime;
return result;
}
@@ -138,6 +164,13 @@ public class VertexLocationHint {
} else if (other.racks != null) {
return false;
}
+ if (containerId != null) {
+ if (!containerId.equals(other.containerId)) {
+ return false;
+ }
+ } else if (other.containerId != null) {
+ return false;
+ }
return true;
}
}
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPluginContext.java
----------------------------------------------------------------------
diff --git a/tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPluginContext.java b/tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPluginContext.java
index bd610c0..5d53b60 100644
--- a/tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPluginContext.java
+++ b/tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPluginContext.java
@@ -22,15 +22,39 @@ import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.hadoop.classification.InterfaceStability.Unstable;
+import org.apache.hadoop.yarn.api.records.Container;
import org.apache.hadoop.yarn.api.records.Resource;
+import org.apache.tez.dag.api.VertexLocationHint.TaskLocationHint;
import org.apache.tez.runtime.api.events.RootInputDataInformationEvent;
+import com.google.common.base.Preconditions;
+
+@Unstable
/**
* Object with API's to interact with the Tez execution engine
*/
public interface VertexManagerPluginContext {
+ public class TaskWithLocationHint {
+ Integer taskIndex;
+ TaskLocationHint locationHint;
+ public TaskWithLocationHint(Integer taskIndex, @Nullable TaskLocationHint locationHint) {
+ Preconditions.checkNotNull(taskIndex);
+ this.taskIndex = taskIndex;
+ this.locationHint = locationHint;
+ }
+
+ public Integer getTaskIndex() {
+ return taskIndex;
+ }
+
+ public TaskLocationHint getTaskLocationHint() {
+ return locationHint;
+ }
+ }
/**
* Get the edge properties on the input edges of this vertex. The input edge
* is represented by the source vertex name
@@ -64,6 +88,13 @@ public interface VertexManagerPluginContext {
Resource getVertexTaskResource();
/**
+ * Get the container for the successful attempt of the task
+ * @return YARN container for the successful task. Maybe null if there is no
+ * successful task.
+ */
+ public Container getTaskContainer(String vertexName, Integer taskIndex);
+
+ /**
* Get the total resource allocated to this vertex. If the DAG is running in
* a busy cluster then it may have no resources available dedicated to it. The
* DAG may divide its available resource among member vertices.
@@ -110,8 +141,8 @@ public interface VertexManagerPluginContext {
* Notify the vertex to start the given tasks
* @param taskIDs Indices of the tasks to be started
*/
- public void scheduleVertexTasks(List<Integer> taskIDs);
-
+ public void scheduleVertexTasks(List<TaskWithLocationHint> tasks);
+
/**
* Get the names of the non-vertex inputs of this vertex. These are primary
* sources of data.
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPluginDescriptor.java
----------------------------------------------------------------------
diff --git a/tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPluginDescriptor.java b/tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPluginDescriptor.java
index 807b7bb..963df48 100644
--- a/tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPluginDescriptor.java
+++ b/tez-api/src/main/java/org/apache/tez/dag/api/VertexManagerPluginDescriptor.java
@@ -27,8 +27,8 @@ public class VertexManagerPluginDescriptor extends TezEntityDescriptor {
super();
}
- public VertexManagerPluginDescriptor(String outputClassName) {
- super(outputClassName);
+ public VertexManagerPluginDescriptor(String className) {
+ super(className);
}
@Override
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-dag/src/main/java/org/apache/tez/dag/app/dag/Task.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/Task.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/Task.java
index bcc9dcd..d2965a9 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/Task.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/Task.java
@@ -40,6 +40,7 @@ public interface Task {
float getProgress();
Map<TezTaskAttemptID, TaskAttempt> getAttempts();
TaskAttempt getAttempt(TezTaskAttemptID attemptID);
+ TaskAttempt getSuccessfulAttempt();
/** Has Task reached the final state or not.
*/
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-dag/src/main/java/org/apache/tez/dag/app/dag/TaskAttempt.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/TaskAttempt.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/TaskAttempt.java
index 2af1232..0563037 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/TaskAttempt.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/TaskAttempt.java
@@ -21,6 +21,7 @@ package org.apache.tez.dag.app.dag;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.hadoop.yarn.api.records.Container;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.api.records.NodeId;
import org.apache.tez.common.counters.DAGCounter;
@@ -83,6 +84,11 @@ public interface TaskAttempt {
* @return the container ID if a container is assigned, otherwise null.
*/
ContainerId getAssignedContainerID();
+
+ /**
+ * @return the container if assigned, otherwise null
+ */
+ Container getAssignedContainer();
/**
* @return container mgr address if a container is assigned, otherwise null.
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-dag/src/main/java/org/apache/tez/dag/app/dag/Vertex.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/Vertex.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/Vertex.java
index a6ccc33..2ff20b8 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/Vertex.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/Vertex.java
@@ -30,6 +30,7 @@ import org.apache.tez.dag.api.OutputDescriptor;
import org.apache.tez.dag.api.ProcessorDescriptor;
import org.apache.tez.dag.api.VertexLocationHint;
import org.apache.tez.dag.api.VertexLocationHint.TaskLocationHint;
+import org.apache.tez.dag.api.VertexManagerPluginContext.TaskWithLocationHint;
import org.apache.tez.dag.api.client.StatusGetOpts;
import org.apache.tez.dag.api.records.DAGProtos.RootInputLeafOutputProto;
import org.apache.tez.dag.api.records.DAGProtos.VertexPlan;
@@ -107,7 +108,7 @@ public interface Vertex extends Comparable<Vertex> {
int getInputVerticesCount();
int getOutputVerticesCount();
- void scheduleTasks(List<Integer> taskIDs);
+ void scheduleTasks(List<TaskWithLocationHint> tasks);
Resource getTaskResource();
ProcessorDescriptor getProcessorDescriptor();
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/ImmediateStartVertexManager.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/ImmediateStartVertexManager.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/ImmediateStartVertexManager.java
index af128ba..cf13f64 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/ImmediateStartVertexManager.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/ImmediateStartVertexManager.java
@@ -18,16 +18,18 @@
package org.apache.tez.dag.app.dag.impl;
-import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import org.apache.tez.dag.api.InputDescriptor;
import org.apache.tez.dag.api.VertexManagerPlugin;
import org.apache.tez.dag.api.VertexManagerPluginContext;
+import org.apache.tez.dag.api.VertexManagerPluginContext.TaskWithLocationHint;
import org.apache.tez.runtime.api.Event;
import org.apache.tez.runtime.api.events.VertexManagerEvent;
+import com.google.common.collect.Lists;
+
/**
* Starts all tasks immediately on vertex start
*/
@@ -41,9 +43,9 @@ public class ImmediateStartVertexManager implements VertexManagerPlugin {
@Override
public void onVertexStarted(Map<String, List<Integer>> completions) {
int numTasks = context.getVertexNumTasks(context.getVertexName());
- List<Integer> scheduledTasks = new ArrayList<Integer>(numTasks);
+ List<TaskWithLocationHint> scheduledTasks = Lists.newArrayListWithCapacity(numTasks);
for (int i=0; i<numTasks; ++i) {
- scheduledTasks.add(new Integer(i));
+ scheduledTasks.add(new TaskWithLocationHint(new Integer(i), null));
}
context.scheduleVertexTasks(scheduledTasks);
}
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/RootInputVertexManager.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/RootInputVertexManager.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/RootInputVertexManager.java
index bbba284..cf68d5e 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/RootInputVertexManager.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/RootInputVertexManager.java
@@ -18,7 +18,6 @@
package org.apache.tez.dag.app.dag.impl;
-import java.util.ArrayList;
import java.util.List;
import java.util.Map;
@@ -26,6 +25,7 @@ import org.apache.tez.dag.api.InputDescriptor;
import org.apache.tez.dag.api.VertexLocationHint;
import org.apache.tez.dag.api.VertexManagerPlugin;
import org.apache.tez.dag.api.VertexManagerPluginContext;
+import org.apache.tez.dag.api.VertexManagerPluginContext.TaskWithLocationHint;
import org.apache.tez.runtime.api.Event;
import org.apache.tez.runtime.api.events.RootInputConfigureVertexTasksEvent;
import org.apache.tez.runtime.api.events.RootInputDataInformationEvent;
@@ -47,9 +47,9 @@ public class RootInputVertexManager implements VertexManagerPlugin {
@Override
public void onVertexStarted(Map<String, List<Integer>> completions) {
int numTasks = context.getVertexNumTasks(context.getVertexName());
- List<Integer> scheduledTasks = new ArrayList<Integer>(numTasks);
+ List<TaskWithLocationHint> scheduledTasks = Lists.newArrayListWithCapacity(numTasks);
for (int i=0; i<numTasks; ++i) {
- scheduledTasks.add(new Integer(i));
+ scheduledTasks.add(new TaskWithLocationHint(new Integer(i), null));
}
context.scheduleVertexTasks(scheduledTasks);
}
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskAttemptImpl.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskAttemptImpl.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskAttemptImpl.java
index bda1d9c..ea3b5b7 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskAttemptImpl.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskAttemptImpl.java
@@ -130,7 +130,8 @@ public class TaskAttemptImpl implements TaskAttempt,
private String trackerName;
private int httpPort;
- // TODO Can these be replaced by the container object ?
+ // TODO Can these be replaced by the container object TEZ-1037
+ private Container container;
private ContainerId containerId;
private NodeId containerNodeId;
private String nodeHttpAddress;
@@ -637,6 +638,16 @@ public class TaskAttemptImpl implements TaskAttempt,
readLock.unlock();
}
}
+
+ @Override
+ public Container getAssignedContainer() {
+ readLock.lock();
+ try {
+ return container;
+ } finally {
+ readLock.unlock();
+ }
+ }
@Override
public String getAssignedContainerMgrAddress() {
@@ -1179,6 +1190,7 @@ public class TaskAttemptImpl implements TaskAttempt,
Container container = ta.appContext.getAllContainers()
.get(event.getContainerId()).getContainer();
+ ta.container = container;
ta.containerId = event.getContainerId();
ta.containerNodeId = container.getNodeId();
ta.nodeHttpAddress = StringInterner.weakIntern(container.getNodeHttpAddress());
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskImpl.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskImpl.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskImpl.java
index 60d5b4f..005088d 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskImpl.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/TaskImpl.java
@@ -368,7 +368,7 @@ public class TaskImpl implements Task, EventHandler<TaskEvent> {
readLock.unlock();
}
}
-
+
@Override
public TaskAttempt getAttempt(TezTaskAttemptID attemptID) {
readLock.lock();
@@ -773,7 +773,8 @@ public class TaskImpl implements Task, EventHandler<TaskEvent> {
(failedAttempts > 0), taskResource, containerContext, leafVertex);
}
- protected TaskAttempt getSuccessfulAttempt() {
+ @Override
+ public TaskAttempt getSuccessfulAttempt() {
readLock.lock();
try {
if (null == successfulAttempt) {
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/VertexImpl.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/VertexImpl.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/VertexImpl.java
index 8ee6980..f20c1a0 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/VertexImpl.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/VertexImpl.java
@@ -62,6 +62,7 @@ import org.apache.tez.dag.api.ProcessorDescriptor;
import org.apache.tez.dag.api.TezUncheckedException;
import org.apache.tez.dag.api.VertexLocationHint;
import org.apache.tez.dag.api.VertexLocationHint.TaskLocationHint;
+import org.apache.tez.dag.api.VertexManagerPluginContext.TaskWithLocationHint;
import org.apache.tez.dag.api.VertexManagerPluginDescriptor;
import org.apache.tez.dag.api.client.ProgressBuilder;
import org.apache.tez.dag.api.client.StatusGetOpts;
@@ -547,7 +548,7 @@ public class VertexImpl implements org.apache.tez.dag.app.dag.Vertex,
private Map<String, VertexGroupInfo> dagVertexGroups;
- private VertexLocationHint vertexLocationHint;
+ private TaskLocationHint taskLocationHints[];
private Map<String, LocalResource> localResources;
private Map<String, String> environment;
private final String javaOpts;
@@ -587,11 +588,11 @@ public class VertexImpl implements org.apache.tez.dag.app.dag.Vertex,
this.readLock = readWriteLock.readLock();
this.writeLock = readWriteLock.writeLock();
- this.vertexLocationHint = vertexLocationHint;
if (LOG.isDebugEnabled()) {
- logLocationHints(this.vertexName, this.vertexLocationHint);
+ logLocationHints(this.vertexName, vertexLocationHint);
}
-
+ setTaskLocationHints(vertexLocationHint);
+
this.dagUgi = appContext.getCurrentDAG().getDagUGI();
this.taskResource = DagTypeConverters
@@ -848,11 +849,11 @@ public class VertexImpl implements org.apache.tez.dag.app.dag.Vertex,
public TaskLocationHint getTaskLocationHint(TezTaskID taskId) {
this.readLock.lock();
try {
- if (vertexLocationHint == null ||
- vertexLocationHint.getTaskLocationHints().size() <= taskId.getId()) {
+ if (taskLocationHints == null ||
+ taskLocationHints.length <= taskId.getId()) {
return null;
}
- return vertexLocationHint.getTaskLocationHints().get(taskId.getId());
+ return taskLocationHints[taskId.getId()];
} finally {
this.readLock.unlock();
}
@@ -972,7 +973,7 @@ public class VertexImpl implements org.apache.tez.dag.app.dag.Vertex,
VertexParallelismUpdatedEvent updatedEvent =
(VertexParallelismUpdatedEvent) historyEvent;
if (updatedEvent.getVertexLocationHint() != null) {
- vertexLocationHint = updatedEvent.getVertexLocationHint();
+ setTaskLocationHints(updatedEvent.getVertexLocationHint());
}
numTasks = updatedEvent.getNumTasks();
handleParallelismUpdate(numTasks, updatedEvent.getSourceEdgeManagers());
@@ -1015,12 +1016,19 @@ public class VertexImpl implements org.apache.tez.dag.app.dag.Vertex,
}
}
+
+ private void setTaskLocationHints(VertexLocationHint vertexLocationHint) {
+ if (vertexLocationHint != null &&
+ vertexLocationHint.getTaskLocationHints() != null &&
+ !vertexLocationHint.getTaskLocationHints().isEmpty()) {
+ List<TaskLocationHint> locHints = vertexLocationHint.getTaskLocationHints();
+ taskLocationHints = locHints.toArray(new TaskLocationHint[locHints.size()]);
+ }
+ }
- // TODO Create InputReadyVertexManager that schedules when there is something
- // to read and use that as default instead of ImmediateStart.TEZ-480
@Override
- public void scheduleTasks(List<Integer> taskIDs) {
- readLock.lock();
+ public void scheduleTasks(List<TaskWithLocationHint> tasksToSchedule) {
+ writeLock.lock();
try {
tasksNotYetScheduled = false;
if (!pendingTaskEvents.isEmpty()) {
@@ -1028,17 +1036,24 @@ public class VertexImpl implements org.apache.tez.dag.app.dag.Vertex,
new VertexEventRouteEvent(getVertexId(), pendingTaskEvents));
pendingTaskEvents.clear();
}
- for (Integer taskID : taskIDs) {
- if (tasks.size() <= taskID.intValue()) {
+ for (TaskWithLocationHint task : tasksToSchedule) {
+ if (numTasks <= task.getTaskIndex().intValue()) {
throw new TezUncheckedException(
- "Invalid taskId: " + taskID + " for vertex: " + vertexName);
+ "Invalid taskId: " + task.getTaskIndex() + " for vertex: " + logIdentifier);
+ }
+ TaskLocationHint locationHint = task.getTaskLocationHint();
+ if (locationHint != null) {
+ if (taskLocationHints == null) {
+ taskLocationHints = new TaskLocationHint[numTasks];
+ }
+ taskLocationHints[task.getTaskIndex().intValue()] = locationHint;
}
eventHandler.handle(new TaskEvent(
- TezTaskID.getInstance(vertexId, taskID.intValue()),
+ TezTaskID.getInstance(vertexId, task.getTaskIndex().intValue()),
TaskEventType.T_SCHEDULE));
}
} finally {
- readLock.unlock();
+ writeLock.unlock();
}
}
@@ -1081,7 +1096,7 @@ public class VertexImpl implements org.apache.tez.dag.app.dag.Vertex,
writeLock.lock();
try {
if (parallelismSet == true) {
- LOG.info("Parallelism can only be set dynamically once per vertex");
+ LOG.info("Parallelism can only be set dynamically once per vertex: " + logIdentifier);
return false;
}
@@ -1120,10 +1135,15 @@ public class VertexImpl implements org.apache.tez.dag.app.dag.Vertex,
return false;
}
if (parallelism == numTasks) {
- LOG.info("setParallelism same as current value: " + parallelism);
+ LOG.info("setParallelism same as current value: " + parallelism +
+ " for vertex: " + logIdentifier);
Preconditions
.checkArgument(sourceEdgeManagers != null,
"Source edge managers must be set when not changing parallelism");
+ } else {
+ LOG.info(
+ "Resetting vertex location hints due to change in parallelism for vertex: " + logIdentifier);
+ vertexLocationHint = null;
}
// start buffering incoming events so that we can re-route existing events
@@ -1215,10 +1235,10 @@ public class VertexImpl implements org.apache.tez.dag.app.dag.Vertex,
public void setVertexLocationHint(VertexLocationHint vertexLocationHint) {
writeLock.lock();
try {
- this.vertexLocationHint = vertexLocationHint;
if (LOG.isDebugEnabled()) {
- logLocationHints(this.vertexName, this.vertexLocationHint);
+ logLocationHints(this.vertexName, vertexLocationHint);
}
+ setTaskLocationHints(vertexLocationHint);
} finally {
writeLock.unlock();
}
@@ -3305,8 +3325,8 @@ public class VertexImpl implements org.apache.tez.dag.app.dag.Vertex,
}
@VisibleForTesting
- VertexLocationHint getVertexLocationHint() {
- return this.vertexLocationHint;
+ TaskLocationHint[] getTaskLocationHints() {
+ return taskLocationHints;
}
// TODO Eventually remove synchronization.
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/VertexManager.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/VertexManager.java b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/VertexManager.java
index b7c3611..a6289ad 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/VertexManager.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/dag/impl/VertexManager.java
@@ -27,6 +27,7 @@ import java.util.Set;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
+import org.apache.hadoop.yarn.api.records.Container;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.tez.common.RuntimeUtils;
import org.apache.tez.common.TezUtils;
@@ -39,6 +40,8 @@ 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.app.AppContext;
+import org.apache.tez.dag.app.dag.Task;
+import org.apache.tez.dag.app.dag.TaskAttempt;
import org.apache.tez.dag.app.dag.Vertex;
import org.apache.tez.dag.app.dag.event.VertexEventRouteEvent;
import org.apache.tez.dag.records.TezTaskAttemptID;
@@ -102,8 +105,8 @@ public class VertexManager {
}
@Override
- public void scheduleVertexTasks(List<Integer> taskIDs) {
- managedVertex.scheduleTasks(taskIDs);
+ public void scheduleVertexTasks(List<TaskWithLocationHint> tasks) {
+ managedVertex.scheduleTasks(tasks);
}
@Override
@@ -181,6 +184,17 @@ public class VertexManager {
public int getNumClusterNodes() {
return appContext.getTaskScheduler().getNumClusterNodes();
}
+
+ @Override
+ public Container getTaskContainer(String vertexName, Integer taskIndex) {
+ Vertex vertex = appContext.getCurrentDAG().getVertex(vertexName);
+ Task task = vertex.getTask(taskIndex.intValue());
+ TaskAttempt attempt = task.getSuccessfulAttempt();
+ if (attempt != null) {
+ return attempt.getAssignedContainer();
+ }
+ return null;
+ }
}
public VertexManager(VertexManagerPlugin plugin,
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerEventHandler.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerEventHandler.java b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerEventHandler.java
index 8fdcfc2..bc1d1b7 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerEventHandler.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerEventHandler.java
@@ -332,11 +332,24 @@ public class TaskSchedulerEventHandler extends AbstractService
event.getCapability().setMemory(mapResourceReqt);*/
TaskAttempt taskAttempt = event.getTaskAttempt();
TaskLocationHint locationHint = event.getLocationHint();
- String hosts[] = (locationHint != null && locationHint.getDataLocalHosts() != null) ? locationHint
- .getDataLocalHosts().toArray(
- new String[locationHint.getDataLocalHosts().size()]) : null;
- String racks[] = (locationHint != null && locationHint.getRacks() != null) ? locationHint
- .getRacks().toArray(new String[locationHint.getRacks().size()]) : null;
+ String hosts[] = null;
+ String racks[] = null;
+ // Until TEZ-1039 is done we need to translate container affinity to the node
+ // for that container
+ if (locationHint != null) {
+ if (locationHint.getAffinitizedContainer() != null) {
+ ContainerId containerId = locationHint.getAffinitizedContainer();
+ AMContainer container = appContext.getAllContainers().get(containerId);
+ hosts = new String[1];
+ hosts[0] = container.getContainer().getNodeId().getHost();
+ } else {
+ hosts = (locationHint.getDataLocalHosts() != null) ? locationHint
+ .getDataLocalHosts().toArray(
+ new String[locationHint.getDataLocalHosts().size()]) : null;
+ racks = (locationHint.getRacks() != null) ? locationHint.getRacks()
+ .toArray(new String[locationHint.getRacks().size()]) : null;
+ }
+ }
taskScheduler.allocateTask(taskAttempt,
event.getCapability(),
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-dag/src/test/java/org/apache/tez/dag/app/dag/impl/TestVertexImpl.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/test/java/org/apache/tez/dag/app/dag/impl/TestVertexImpl.java b/tez-dag/src/test/java/org/apache/tez/dag/app/dag/impl/TestVertexImpl.java
index aa829df..518a3c5 100644
--- a/tez-dag/src/test/java/org/apache/tez/dag/app/dag/impl/TestVertexImpl.java
+++ b/tez-dag/src/test/java/org/apache/tez/dag/app/dag/impl/TestVertexImpl.java
@@ -37,7 +37,6 @@ import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
-import java.util.Map.Entry;
import com.google.protobuf.ByteString;
import org.apache.commons.logging.Log;
@@ -56,13 +55,13 @@ import org.apache.hadoop.yarn.util.Clock;
import org.apache.hadoop.yarn.util.SystemClock;
import org.apache.tez.dag.api.DagTypeConverters;
import org.apache.tez.dag.api.EdgeManager;
-import org.apache.tez.dag.api.EdgeManagerContext;
import org.apache.tez.dag.api.EdgeManagerDescriptor;
import org.apache.tez.dag.api.EdgeProperty;
import org.apache.tez.dag.api.InputDescriptor;
import org.apache.tez.dag.api.TezConfiguration;
import org.apache.tez.dag.api.VertexLocationHint;
import org.apache.tez.dag.api.VertexLocationHint.TaskLocationHint;
+import org.apache.tez.dag.api.VertexManagerPluginContext.TaskWithLocationHint;
import org.apache.tez.dag.api.client.VertexStatus;
import org.apache.tez.dag.api.oldrecords.TaskState;
import org.apache.tez.dag.api.records.DAGProtos;
@@ -117,7 +116,6 @@ import org.apache.tez.runtime.api.OutputCommitter;
import org.apache.tez.runtime.api.OutputCommitterContext;
import org.apache.tez.runtime.api.events.CompositeDataMovementEvent;
import org.apache.tez.runtime.api.events.DataMovementEvent;
-import org.apache.tez.runtime.api.events.InputReadErrorEvent;
import org.apache.tez.runtime.api.events.RootInputConfigureVertexTasksEvent;
import org.apache.tez.runtime.api.events.RootInputDataInformationEvent;
import org.apache.tez.test.EdgeManagerForTest;
@@ -1455,7 +1453,7 @@ public class TestVertexImpl {
new VertexEventRouteEvent(v3.getVertexId(), taskEvents));
dispatcher.await();
Assert.assertEquals(2, v3.pendingTaskEvents.size());
- v3.scheduleTasks(Collections.singletonList(new Integer(0)));
+ v3.scheduleTasks(Collections.singletonList(new TaskWithLocationHint(new Integer(0), null)));
dispatcher.await();
Assert.assertEquals(0, v3.pendingTaskEvents.size());
// send events and test that they are not buffered anymore
@@ -2112,7 +2110,9 @@ public class TestVertexImpl {
Assert.assertEquals(numTasks, v1.getTotalTasks());
Assert.assertEquals(RootInputVertexManager.class.getName(), v1
.getVertexManager().getPlugin().getClass().getName());
- Assert.assertEquals(v1Hints, v1.getVertexLocationHint().getTaskLocationHints());
+ for (int i=0; i < v1Hints.size(); ++i) {
+ Assert.assertEquals(v1Hints.get(i), v1.getTaskLocationHints()[i]);
+ }
Assert.assertEquals(true, runner1.hasShutDown);
Assert.assertEquals(numTasks, vertices.get("vertex2").getTotalTasks());
@@ -2201,7 +2201,9 @@ public class TestVertexImpl {
Assert.assertEquals(5, v1.getTotalTasks());
Assert.assertEquals(RootInputVertexManager.class.getName(), v1
.getVertexManager().getPlugin().getClass().getName());
- Assert.assertEquals(v1Hints, v1.getVertexLocationHint().getTaskLocationHints());
+ for (int i=0; i < v1Hints.size(); ++i) {
+ Assert.assertEquals(v1Hints.get(i), v1.getTaskLocationHints()[i]);
+ }
Assert.assertEquals(true, runner1.hasShutDown);
VertexImplWithCustomInitializer v2 = (VertexImplWithCustomInitializer) vertices.get("vertex2");
@@ -2214,7 +2216,9 @@ public class TestVertexImpl {
Assert.assertEquals(10, v2.getTotalTasks());
Assert.assertEquals(RootInputVertexManager.class.getName(), v2
.getVertexManager().getPlugin().getClass().getName());
- Assert.assertEquals(v2Hints, v2.getVertexLocationHint().getTaskLocationHints());
+ for (int i=0; i < v2Hints.size(); ++i) {
+ Assert.assertEquals(v2Hints.get(i), v2.getTaskLocationHints()[i]);
+ }
Assert.assertEquals(true, runner2.hasShutDown);
}
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-mapreduce-examples/src/main/java/org/apache/tez/mapreduce/examples/BroadcastAndOneToOneExample.java
----------------------------------------------------------------------
diff --git a/tez-mapreduce-examples/src/main/java/org/apache/tez/mapreduce/examples/BroadcastAndOneToOneExample.java b/tez-mapreduce-examples/src/main/java/org/apache/tez/mapreduce/examples/BroadcastAndOneToOneExample.java
new file mode 100644
index 0000000..db3e1cf
--- /dev/null
+++ b/tez-mapreduce-examples/src/main/java/org/apache/tez/mapreduce/examples/BroadcastAndOneToOneExample.java
@@ -0,0 +1,287 @@
+/**
+ * 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.mapreduce.examples;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.io.IntWritable;
+import org.apache.hadoop.io.Text;
+import org.apache.hadoop.mapred.JobConf;
+import org.apache.hadoop.mapreduce.security.TokenCache;
+import org.apache.hadoop.security.Credentials;
+import org.apache.hadoop.security.UserGroupInformation;
+import org.apache.hadoop.yarn.api.records.ApplicationId;
+import org.apache.tez.client.AMConfiguration;
+import org.apache.tez.client.TezClient;
+import org.apache.tez.client.TezClientUtils;
+import org.apache.tez.client.TezSession;
+import org.apache.tez.client.TezSessionConfiguration;
+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.InputDescriptor;
+import org.apache.tez.dag.api.OutputDescriptor;
+import org.apache.tez.dag.api.ProcessorDescriptor;
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.api.TezUncheckedException;
+import org.apache.tez.dag.api.Vertex;
+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.VertexManagerPluginDescriptor;
+import org.apache.tez.dag.api.client.DAGClient;
+import org.apache.tez.dag.api.client.DAGStatus;
+import org.apache.tez.dag.library.vertexmanager.InputReadyVertexManager;
+import org.apache.tez.mapreduce.hadoop.MRHelpers;
+import org.apache.tez.mapreduce.hadoop.MRJobConfig;
+import org.apache.tez.mapreduce.hadoop.MultiStageMRConfToTezTranslator;
+import org.apache.tez.runtime.api.AbstractLogicalIOProcessor;
+import org.apache.tez.runtime.api.Event;
+import org.apache.tez.runtime.api.LogicalInput;
+import org.apache.tez.runtime.api.LogicalOutput;
+import org.apache.tez.runtime.library.api.KeyValueReader;
+import org.apache.tez.runtime.library.api.KeyValueWriter;
+import org.apache.tez.runtime.library.input.ShuffledUnorderedKVInput;
+import org.apache.tez.runtime.library.output.OnFileUnorderedKVOutput;
+
+import com.google.common.base.Preconditions;
+
+public class BroadcastAndOneToOneExample {
+ public static class InputProcessor extends AbstractLogicalIOProcessor {
+ Text word = new Text();
+
+ @Override
+ public void initialize() throws Exception {
+ }
+
+ @Override
+ public void handleEvents(List<Event> processorEvents) {
+ }
+
+ @Override
+ public void close() throws Exception {
+ }
+
+ @Override
+ public void run(Map<String, LogicalInput> inputs,
+ Map<String, LogicalOutput> outputs) throws Exception {
+ for (LogicalOutput output : outputs.values()) {
+ output.start();
+ }
+ Preconditions.checkArgument(outputs.size() == 1);
+ OnFileUnorderedKVOutput output = (OnFileUnorderedKVOutput) outputs.values().iterator().next();
+ KeyValueWriter kvWriter = (KeyValueWriter) output.getWriter();
+ kvWriter.write(word, new IntWritable(getContext().getTaskIndex()));
+ }
+ }
+
+ public static class OneToOneProcessor extends AbstractLogicalIOProcessor {
+ Text word = new Text();
+
+ @Override
+ public void initialize() throws Exception {
+ }
+
+ @Override
+ public void handleEvents(List<Event> processorEvents) {
+ }
+
+ @Override
+ public void close() throws Exception {
+ }
+
+ @Override
+ public void run(Map<String, LogicalInput> inputs,
+ Map<String, LogicalOutput> outputs) throws Exception {
+ Preconditions.checkArgument(inputs.size() == 2);
+
+ for (LogicalInput input : inputs.values()) {
+ input.start();
+ }
+
+ KeyValueReader inputKvReader = (KeyValueReader) inputs.get("Input").getReader();
+ KeyValueReader broadcastKvReader = (KeyValueReader) inputs.get("Broadcast").getReader();
+ int sum = 0;
+ while (broadcastKvReader.next()) {
+ sum += ((IntWritable) broadcastKvReader.getCurrentValue()).get();
+ }
+ while (inputKvReader.next()) {
+ sum += ((IntWritable) inputKvReader.getCurrentValue()).get();
+ }
+ System.out.println("Index: " + getContext().getTaskIndex() + " sum: " + sum);
+ int taskIndex = getContext().getTaskIndex();
+ switch (taskIndex) {
+ case 0:
+ Preconditions.checkState((sum==1), "Sum = " + sum);
+ break;
+ case 1:
+ Preconditions.checkState((sum==2), "Sum = " + sum);
+ break;
+ case 2:
+ Preconditions.checkState((sum==3), "Sum = " + sum);
+ break;
+ default:
+ throw new TezUncheckedException("Unexpected taskIndex: " + taskIndex);
+ }
+ }
+
+ }
+
+ private DAG createDAG(FileSystem fs, TezConfiguration tezConf,
+ Path stagingDir) throws IOException {
+ Configuration kvInputConf = new JobConf((Configuration)tezConf);
+ kvInputConf.set(MRJobConfig.MAP_OUTPUT_KEY_CLASS,
+ Text.class.getName());
+ kvInputConf.set(MRJobConfig.MAP_OUTPUT_VALUE_CLASS,
+ IntWritable.class.getName());
+ MultiStageMRConfToTezTranslator.translateVertexConfToTez(kvInputConf,
+ null);
+
+ Configuration kvOneToOneConf = new JobConf((Configuration)tezConf);
+
+ MultiStageMRConfToTezTranslator.translateVertexConfToTez(kvOneToOneConf,
+ kvInputConf);
+
+ MRHelpers.doJobClientMagic(kvInputConf);
+ MRHelpers.doJobClientMagic(kvOneToOneConf);
+
+ byte[] kvInputPayload = MRHelpers.createUserPayloadFromConf(kvInputConf);
+ Vertex broadcastVertex = new Vertex("Broadcast", new ProcessorDescriptor(
+ InputProcessor.class.getName()),
+ 2, MRHelpers.getMapResource(kvInputConf));
+ broadcastVertex.setJavaOpts(MRHelpers.getMapJavaOpts(kvInputConf));
+
+ int numOneToOneTasks = 3;
+ Vertex inputVertex = new Vertex("Input", new ProcessorDescriptor(
+ InputProcessor.class.getName()),
+ numOneToOneTasks, MRHelpers.getMapResource(kvInputConf));
+ inputVertex.setJavaOpts(MRHelpers.getMapJavaOpts(kvInputConf));
+
+ byte[] kvOneToOnePayload = MRHelpers.createUserPayloadFromConf(kvOneToOneConf);
+ Vertex oneToOneVertex = new Vertex("OneToOne",
+ new ProcessorDescriptor(
+ OneToOneProcessor.class.getName()),
+ numOneToOneTasks, MRHelpers.getReduceResource(kvOneToOneConf));
+ oneToOneVertex.setJavaOpts(
+ MRHelpers.getReduceJavaOpts(kvOneToOneConf)).setVertexManagerPlugin(
+ new VertexManagerPluginDescriptor(InputReadyVertexManager.class.getName()));
+
+ DAG dag = new DAG("BroadcastAndOneToOneExample");
+ dag.addVertex(inputVertex)
+ .addVertex(broadcastVertex)
+ .addVertex(oneToOneVertex)
+ .addEdge(
+ new Edge(inputVertex, oneToOneVertex, new EdgeProperty(
+ DataMovementType.ONE_TO_ONE, DataSourceType.PERSISTED,
+ SchedulingType.SEQUENTIAL,
+ new OutputDescriptor(OnFileUnorderedKVOutput.class.getName())
+ .setUserPayload(kvInputPayload),
+ new InputDescriptor(ShuffledUnorderedKVInput.class.getName())
+ .setUserPayload(kvOneToOnePayload))))
+ .addEdge(
+ new Edge(broadcastVertex, oneToOneVertex, new EdgeProperty(
+ DataMovementType.BROADCAST, DataSourceType.PERSISTED,
+ SchedulingType.SEQUENTIAL,
+ new OutputDescriptor(OnFileUnorderedKVOutput.class.getName())
+ .setUserPayload(kvInputPayload),
+ new InputDescriptor(ShuffledUnorderedKVInput.class.getName())
+ .setUserPayload(kvOneToOnePayload))));
+ return dag;
+ }
+
+ private Credentials credentials = new Credentials();
+
+ public boolean run(Configuration conf) throws Exception {
+ System.out.println("Running BroadcastAndOneToOneExample");
+ // conf and UGI
+ TezConfiguration tezConf;
+ if (conf != null) {
+ tezConf = new TezConfiguration(conf);
+ } else {
+ tezConf = new TezConfiguration();
+ }
+ UserGroupInformation.setConfiguration(tezConf);
+ String user = UserGroupInformation.getCurrentUser().getShortUserName();
+
+ TezClient tezClient = new TezClient(tezConf);
+ ApplicationId appId = tezClient.createApplication();
+
+ // staging dir
+ FileSystem fs = FileSystem.get(tezConf);
+ String stagingDirStr = Path.SEPARATOR + "user" + Path.SEPARATOR
+ + user + Path.SEPARATOR+ ".staging" + Path.SEPARATOR
+ + Path.SEPARATOR + appId.toString();
+ Path stagingDir = new Path(stagingDirStr);
+ tezConf.set(TezConfiguration.TEZ_AM_STAGING_DIR, stagingDirStr);
+ stagingDir = fs.makeQualified(stagingDir);
+
+ // security
+ TokenCache.obtainTokensForNamenodes(credentials, new Path[] {stagingDir}, tezConf);
+ TezClientUtils.ensureStagingDirExists(tezConf, stagingDir);
+
+ tezConf.set(TezConfiguration.TEZ_AM_JAVA_OPTS,
+ MRHelpers.getMRAMJavaOpts(tezConf));
+
+ // No need to add jar containing this class as assumed to be part of
+ // the tez jars.
+
+ // TEZ-674 Obtain tokens based on the Input / Output paths. For now assuming staging dir
+ // is the same filesystem as the one used for Input/Output.
+ TezSession tezSession = null;
+ AMConfiguration amConfig = new AMConfiguration(null,
+ null, tezConf, credentials);
+
+ TezSessionConfiguration sessionConfig =
+ new TezSessionConfiguration(amConfig, tezConf);
+ tezSession = new TezSession("WordCountSession", appId,
+ sessionConfig);
+ tezSession.start();
+
+ DAGClient dagClient = null;
+
+ try {
+ DAG dag = createDAG(fs, tezConf, stagingDir);
+
+ tezSession.waitTillReady();
+ dagClient = tezSession.submitDAG(dag);
+
+ // monitoring
+ DAGStatus dagStatus = dagClient.waitForCompletionWithAllStatusUpdates(null);
+ if (dagStatus.getState() != DAGStatus.State.SUCCEEDED) {
+ System.out.println("DAG diagnostics: " + dagStatus.getDiagnostics());
+ return false;
+ }
+ return true;
+ } finally {
+ fs.delete(stagingDir, true);
+ tezSession.stop();
+ }
+ }
+
+ public static void main(String[] args) throws Exception {
+ Configuration conf = new Configuration();
+ BroadcastAndOneToOneExample job = new BroadcastAndOneToOneExample();
+ job.run(conf);
+ }
+}
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/InputReadyVertexManager.java
----------------------------------------------------------------------
diff --git a/tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/InputReadyVertexManager.java b/tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/InputReadyVertexManager.java
index 95e34d7..8e4b109 100644
--- a/tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/InputReadyVertexManager.java
+++ b/tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/InputReadyVertexManager.java
@@ -23,12 +23,15 @@ import java.util.Map;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
+import org.apache.hadoop.yarn.api.records.Container;
import org.apache.tez.dag.api.EdgeProperty;
import org.apache.tez.dag.api.EdgeProperty.DataMovementType;
import org.apache.tez.dag.api.InputDescriptor;
import org.apache.tez.dag.api.TezUncheckedException;
+import org.apache.tez.dag.api.VertexLocationHint.TaskLocationHint;
import org.apache.tez.dag.api.VertexManagerPlugin;
import org.apache.tez.dag.api.VertexManagerPluginContext;
+import org.apache.tez.dag.api.VertexManagerPluginContext.TaskWithLocationHint;
import org.apache.tez.runtime.api.Event;
import org.apache.tez.runtime.api.events.VertexManagerEvent;
@@ -42,7 +45,8 @@ public class InputReadyVertexManager implements VertexManagerPlugin {
VertexManagerPluginContext context;
Map<String, SourceVertexInfo> srcVertexInfo = Maps.newHashMap();
boolean taskIsStarted[];
- int templateOneToOne[];
+ int oneToOneSrcTasksDoneCount[];
+ Container oneToOneLocationHints[];
int numOneToOneEdges;
class SourceVertexInfo {
@@ -67,7 +71,7 @@ public class InputReadyVertexManager implements VertexManagerPlugin {
@Override
public void onVertexStarted(Map<String, List<Integer>> completions) {
int numManagedTasks = context.getVertexNumTasks(context.getVertexName());
- LOG.info("Managing " + numManagedTasks + " for vertex: " + context.getVertexName());
+ LOG.info("Managing " + numManagedTasks + " tasks for vertex: " + context.getVertexName());
taskIsStarted = new boolean[numManagedTasks];
// find out about all input edge types. If there is a custom edge then
@@ -106,19 +110,20 @@ public class InputReadyVertexManager implements VertexManagerPlugin {
throw new TezUncheckedException(
"Managed task number must equal 1-1 source task number");
}
- templateOneToOne = new int[oneToOneSrcTaskCount];
+ oneToOneSrcTasksDoneCount = new int[oneToOneSrcTaskCount];
+ oneToOneLocationHints = new Container[oneToOneSrcTaskCount];
}
for (Map.Entry<String, List<Integer>> entry : completions.entrySet()) {
for (Integer task : entry.getValue()) {
- handleSouceTaskFinished(entry.getKey(), task);
+ handleSourceTaskFinished(entry.getKey(), task);
}
}
}
@Override
public void onSourceTaskCompleted(String srcVertexName, Integer taskId) {
- handleSouceTaskFinished(srcVertexName, taskId);
+ handleSourceTaskFinished(srcVertexName, taskId);
}
@Override
@@ -130,19 +135,23 @@ public class InputReadyVertexManager implements VertexManagerPlugin {
InputDescriptor inputDescriptor, List<Event> events) {
}
- void handleSouceTaskFinished(String vertex, Integer taskId) {
+ void handleSourceTaskFinished(String vertex, Integer taskId) {
SourceVertexInfo srcInfo = srcVertexInfo.get(vertex);
if (srcInfo.taskIsFinished[taskId.intValue()] == null) {
// not a duplicate completion
srcInfo.taskIsFinished[taskId.intValue()] = new Boolean(true);
srcInfo.numFinishedTasks++;
if (srcInfo.edgeProperty.getDataMovementType() == DataMovementType.ONE_TO_ONE) {
- templateOneToOne[taskId.intValue()]++;
+ oneToOneSrcTasksDoneCount[taskId.intValue()]++;
+ // keep the latest container that completed as the location hint
+ // After there is standard data size info available then use it
+ oneToOneLocationHints[taskId.intValue()] = context.getTaskContainer(vertex, taskId);
}
}
// custom edge needs to tell us which of our tasks its connected to
// for now only-built in edges supported
+ // Check if current source task's vertex is completed.
if (srcInfo.edgeProperty.getDataMovementType() != DataMovementType.ONE_TO_ONE
&& srcInfo.numTasks != srcInfo.numFinishedTasks) {
// we depend on all tasks to finish. So nothing to do now.
@@ -161,27 +170,35 @@ public class InputReadyVertexManager implements VertexManagerPlugin {
}
// all source vertices will full dependencies are done
- List<Integer> tasksToStart = null;
+ List<TaskWithLocationHint> tasksToStart = null;
if (numOneToOneEdges == 0) {
// no 1-1 dependency. Start all tasks
- tasksToStart = Lists.newArrayListWithCapacity(taskIsStarted.length);
- for (int i=0; i<taskIsStarted.length; ++i) {
+ int numTasks = taskIsStarted.length;
+ LOG.info("Starting all " + numTasks + "tasks for vertex: " + context.getVertexName());
+ tasksToStart = Lists.newArrayListWithCapacity(numTasks);
+ for (int i=0; i<numTasks; ++i) {
taskIsStarted[i] = true;
- tasksToStart.add(new Integer(i));
+ tasksToStart.add(new TaskWithLocationHint(new Integer(i), null));
}
} else {
// start only the ready 1-1 tasks
tasksToStart = Lists.newLinkedList();
for (int i=0; i<taskIsStarted.length; ++i) {
- if (!taskIsStarted[i] && templateOneToOne[i] == numOneToOneEdges) {
+ if (!taskIsStarted[i] && oneToOneSrcTasksDoneCount[i] == numOneToOneEdges) {
taskIsStarted[i] = true;
- tasksToStart.add(new Integer(i));
+ TaskLocationHint locationHint = null;
+ if (oneToOneLocationHints[i] != null) {
+ locationHint = new TaskLocationHint(oneToOneLocationHints[i].getId());
+ }
+ LOG.info("Starting task " + i + " for vertex: "
+ + context.getVertexName() + " with location: "
+ + ((locationHint != null) ? locationHint.getAffinitizedContainer() : "null"));
+ tasksToStart.add(new TaskWithLocationHint(new Integer(i), locationHint));
}
}
}
if (tasksToStart != null && !tasksToStart.isEmpty()) {
- // TODO determine placement after TEZ-1018
context.scheduleVertexTasks(tasksToStart);
}
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/ShuffleVertexManager.java
----------------------------------------------------------------------
diff --git a/tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/ShuffleVertexManager.java b/tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/ShuffleVertexManager.java
index e0e19c9..d17b367 100644
--- a/tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/ShuffleVertexManager.java
+++ b/tez-runtime-library/src/main/java/org/apache/tez/dag/library/vertexmanager/ShuffleVertexManager.java
@@ -19,7 +19,6 @@
package org.apache.tez.dag.library.vertexmanager;
import java.io.IOException;
-import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
@@ -40,6 +39,7 @@ import org.apache.tez.dag.api.TezUncheckedException;
import org.apache.tez.dag.api.VertexManagerPlugin;
import org.apache.tez.dag.api.VertexManagerPluginContext;
import org.apache.tez.dag.api.EdgeProperty.DataMovementType;
+import org.apache.tez.dag.api.VertexManagerPluginContext.TaskWithLocationHint;
import org.apache.tez.runtime.api.Event;
import org.apache.tez.runtime.api.events.DataMovementEvent;
import org.apache.tez.runtime.api.events.InputReadErrorEvent;
@@ -116,7 +116,7 @@ public class ShuffleVertexManager implements VertexManagerPlugin {
int numSourceTasks = 0;
int numSourceTasksCompleted = 0;
int numVertexManagerEventsReceived = 0;
- ArrayList<Integer> pendingTasks;
+ List<Integer> pendingTasks;
int totalTasksToSchedule = 0;
Map<String, Set<Integer>> bipartiteSources = Maps.newHashMap();
@@ -289,7 +289,7 @@ public class ShuffleVertexManager implements VertexManagerPlugin {
@Override
public void onVertexStarted(Map<String, List<Integer>> completions) {
- pendingTasks = new ArrayList<Integer>(
+ pendingTasks = Lists.newArrayListWithCapacity(
context.getVertexNumTasks(context.getVertexName()));
// track the tasks in this vertex
updatePendingTasks();
@@ -445,10 +445,10 @@ public class ShuffleVertexManager implements VertexManagerPlugin {
parallelismDetermined = true;
determineParallelismAndApply();
}
- ArrayList<Integer> scheduledTasks = new ArrayList<Integer>(numTasksToSchedule);
+ List<TaskWithLocationHint> scheduledTasks = Lists.newArrayListWithCapacity(numTasksToSchedule);
while(!pendingTasks.isEmpty() && numTasksToSchedule > 0) {
numTasksToSchedule--;
- scheduledTasks.add(pendingTasks.get(0));
+ scheduledTasks.add(new TaskWithLocationHint(pendingTasks.get(0), null));
pendingTasks.remove(0);
}
context.scheduleVertexTasks(scheduledTasks);
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-runtime-library/src/test/java/org/apache/tez/dag/library/vertexmanager/TestInputReadyVertexManager.java
----------------------------------------------------------------------
diff --git a/tez-runtime-library/src/test/java/org/apache/tez/dag/library/vertexmanager/TestInputReadyVertexManager.java b/tez-runtime-library/src/test/java/org/apache/tez/dag/library/vertexmanager/TestInputReadyVertexManager.java
index 289f004..65e63d6 100644
--- a/tez-runtime-library/src/test/java/org/apache/tez/dag/library/vertexmanager/TestInputReadyVertexManager.java
+++ b/tez-runtime-library/src/test/java/org/apache/tez/dag/library/vertexmanager/TestInputReadyVertexManager.java
@@ -25,12 +25,15 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import org.apache.hadoop.yarn.api.records.Container;
+import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.tez.dag.api.EdgeProperty;
import org.apache.tez.dag.api.InputDescriptor;
import org.apache.tez.dag.api.OutputDescriptor;
import org.apache.tez.dag.api.TezUncheckedException;
import org.apache.tez.dag.api.VertexManagerPluginContext;
import org.apache.tez.dag.api.EdgeProperty.SchedulingType;
+import org.apache.tez.dag.api.VertexManagerPluginContext.TaskWithLocationHint;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
@@ -44,7 +47,7 @@ import com.google.common.collect.Maps;
public class TestInputReadyVertexManager {
@Captor
- ArgumentCaptor<List<Integer>> requestCaptor;
+ ArgumentCaptor<List<TaskWithLocationHint>> requestCaptor;
@Before
public void init() {
@@ -98,12 +101,21 @@ public class TestInputReadyVertexManager {
new InputDescriptor("in"));
String mockManagedVertexId = "Vertex";
+ Container mockContainer1 = mock(Container.class);
+ ContainerId mockCId1 = mock(ContainerId.class);
+ when(mockContainer1.getId()).thenReturn(mockCId1);
+ Container mockContainer2 = mock(Container.class);
+ ContainerId mockCId2 = mock(ContainerId.class);
+ when(mockContainer2.getId()).thenReturn(mockCId2);
VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
when(mockContext.getVertexName()).thenReturn(mockManagedVertexId);
when(mockContext.getVertexNumTasks(mockManagedVertexId)).thenReturn(3);
when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(3);
+ when(mockContext.getTaskContainer(mockSrcVertexId1, 0)).thenReturn(mockContainer1);
+ when(mockContext.getTaskContainer(mockSrcVertexId1, 1)).thenReturn(mockContainer2);
+ when(mockContext.getTaskContainer(mockSrcVertexId1, 2)).thenReturn(mockContainer1);
mockInputVertices.put(mockSrcVertexId1, eProp1);
Map<String, List<Integer>> initialCompletions = Maps.newHashMap();
@@ -114,15 +126,21 @@ public class TestInputReadyVertexManager {
manager.onVertexStarted(initialCompletions);
verify(mockContext, times(1)).scheduleVertexTasks(requestCaptor.capture());
Assert.assertEquals(1, requestCaptor.getValue().size());
- Assert.assertEquals(0, requestCaptor.getValue().get(0).intValue());
+ Assert.assertEquals(0, requestCaptor.getValue().get(0).getTaskIndex().intValue());
+ Assert.assertEquals(mockCId1, requestCaptor.getValue().get(0)
+ .getTaskLocationHint().getAffinitizedContainer());
manager.onSourceTaskCompleted(mockSrcVertexId1, 1);
verify(mockContext, times(2)).scheduleVertexTasks(requestCaptor.capture());
Assert.assertEquals(1, requestCaptor.getValue().size());
- Assert.assertEquals(1, requestCaptor.getValue().get(0).intValue());
+ Assert.assertEquals(1, requestCaptor.getValue().get(0).getTaskIndex().intValue());
+ Assert.assertEquals(mockCId2, requestCaptor.getValue().get(0)
+ .getTaskLocationHint().getAffinitizedContainer());
manager.onSourceTaskCompleted(mockSrcVertexId1, 2);
verify(mockContext, times(3)).scheduleVertexTasks(requestCaptor.capture());
Assert.assertEquals(1, requestCaptor.getValue().size());
- Assert.assertEquals(2, requestCaptor.getValue().get(0).intValue());
+ Assert.assertEquals(2, requestCaptor.getValue().get(0).getTaskIndex().intValue());
+ Assert.assertEquals(mockCId1, requestCaptor.getValue().get(0)
+ .getTaskLocationHint().getAffinitizedContainer());
}
@Test (timeout=5000)
@@ -152,6 +170,12 @@ public class TestInputReadyVertexManager {
new InputDescriptor("in"));
String mockManagedVertexId = "Vertex";
+ Container mockContainer2 = mock(Container.class);
+ ContainerId mockCId2 = mock(ContainerId.class);
+ when(mockContainer2.getId()).thenReturn(mockCId2);
+ Container mockContainer3 = mock(Container.class);
+ ContainerId mockCId3 = mock(ContainerId.class);
+ when(mockContainer3.getId()).thenReturn(mockCId3);
VertexManagerPluginContext mockContext = mock(VertexManagerPluginContext.class);
when(mockContext.getInputVertexEdgeProperties()).thenReturn(mockInputVertices);
@@ -159,6 +183,12 @@ public class TestInputReadyVertexManager {
when(mockContext.getVertexNumTasks(mockSrcVertexId1)).thenReturn(3);
when(mockContext.getVertexNumTasks(mockSrcVertexId2)).thenReturn(3);
when(mockContext.getVertexNumTasks(mockSrcVertexId3)).thenReturn(3);
+ when(mockContext.getTaskContainer(mockSrcVertexId2, 0)).thenReturn(mockContainer2);
+ when(mockContext.getTaskContainer(mockSrcVertexId2, 1)).thenReturn(mockContainer2);
+ when(mockContext.getTaskContainer(mockSrcVertexId2, 2)).thenReturn(mockContainer2);
+ when(mockContext.getTaskContainer(mockSrcVertexId3, 0)).thenReturn(mockContainer3);
+ when(mockContext.getTaskContainer(mockSrcVertexId3, 1)).thenReturn(mockContainer3);
+ when(mockContext.getTaskContainer(mockSrcVertexId3, 2)).thenReturn(mockContainer3);
mockInputVertices.put(mockSrcVertexId1, eProp1);
mockInputVertices.put(mockSrcVertexId2, eProp2);
mockInputVertices.put(mockSrcVertexId3, eProp3);
@@ -203,20 +233,26 @@ public class TestInputReadyVertexManager {
manager.onSourceTaskCompleted(mockSrcVertexId1, 2); // v1 done
verify(mockContext, times(1)).scheduleVertexTasks(requestCaptor.capture());
Assert.assertEquals(1, requestCaptor.getValue().size());
- Assert.assertEquals(0, requestCaptor.getValue().get(0).intValue());
+ Assert.assertEquals(0, requestCaptor.getValue().get(0).getTaskIndex().intValue());
+ Assert.assertEquals(mockCId3, requestCaptor.getValue().get(0)
+ .getTaskLocationHint().getAffinitizedContainer()); // affinity to last completion
// 1-1 completion triggers since other 1-1 is done
manager.onSourceTaskCompleted(mockSrcVertexId3, 1);
verify(mockContext, times(2)).scheduleVertexTasks(requestCaptor.capture());
Assert.assertEquals(1, requestCaptor.getValue().size());
- Assert.assertEquals(1, requestCaptor.getValue().get(0).intValue());
+ Assert.assertEquals(1, requestCaptor.getValue().get(0).getTaskIndex().intValue());
+ Assert.assertEquals(mockCId3, requestCaptor.getValue().get(0)
+ .getTaskLocationHint().getAffinitizedContainer()); // affinity to last completion
// 1-1 completion does not trigger since other 1-1 is not done
- manager.onSourceTaskCompleted(mockSrcVertexId2, 2);
+ manager.onSourceTaskCompleted(mockSrcVertexId3, 2);
verify(mockContext, times(2)).scheduleVertexTasks(anyList());
// 1-1 completion trigger start
- manager.onSourceTaskCompleted(mockSrcVertexId3, 2);
+ manager.onSourceTaskCompleted(mockSrcVertexId2, 2);
verify(mockContext, times(3)).scheduleVertexTasks(requestCaptor.capture());
Assert.assertEquals(1, requestCaptor.getValue().size());
- Assert.assertEquals(2, requestCaptor.getValue().get(0).intValue());
+ Assert.assertEquals(2, requestCaptor.getValue().get(0).getTaskIndex().intValue());
+ Assert.assertEquals(mockCId2, requestCaptor.getValue().get(0)
+ .getTaskLocationHint().getAffinitizedContainer()); // affinity to last completion
// no more starts
manager.onSourceTaskCompleted(mockSrcVertexId3, 2);
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-runtime-library/src/test/java/org/apache/tez/dag/library/vertexmanager/TestShuffleVertexManager.java
----------------------------------------------------------------------
diff --git a/tez-runtime-library/src/test/java/org/apache/tez/dag/library/vertexmanager/TestShuffleVertexManager.java b/tez-runtime-library/src/test/java/org/apache/tez/dag/library/vertexmanager/TestShuffleVertexManager.java
index 844cb1a..99f4245 100644
--- a/tez-runtime-library/src/test/java/org/apache/tez/dag/library/vertexmanager/TestShuffleVertexManager.java
+++ b/tez-runtime-library/src/test/java/org/apache/tez/dag/library/vertexmanager/TestShuffleVertexManager.java
@@ -38,6 +38,7 @@ import org.apache.tez.dag.api.TezUncheckedException;
import org.apache.tez.dag.api.VertexLocationHint;
import org.apache.tez.dag.api.EdgeProperty.SchedulingType;
import org.apache.tez.dag.api.VertexManagerPluginContext;
+import org.apache.tez.dag.api.VertexManagerPluginContext.TaskWithLocationHint;
import org.apache.tez.runtime.api.events.DataMovementEvent;
import org.apache.tez.runtime.api.events.VertexManagerEvent;
import org.apache.tez.runtime.library.shuffle.impl.ShuffleUserPayloads.VertexManagerEventPayloadProto;
@@ -108,7 +109,10 @@ public class TestShuffleVertexManager {
public Object answer(InvocationOnMock invocation) {
Object[] args = invocation.getArguments();
scheduledTasks.clear();
- scheduledTasks.addAll((List<Integer>)args[0]);
+ List<TaskWithLocationHint> tasks = (List<TaskWithLocationHint>)args[0];
+ for (TaskWithLocationHint task : tasks) {
+ scheduledTasks.add(task.getTaskIndex());
+ }
return null;
}}).when(mockContext).scheduleVertexTasks(anyList());
@@ -312,7 +316,10 @@ public class TestShuffleVertexManager {
public Object answer(InvocationOnMock invocation) {
Object[] args = invocation.getArguments();
scheduledTasks.clear();
- scheduledTasks.addAll((List<Integer>)args[0]);
+ List<TaskWithLocationHint> tasks = (List<TaskWithLocationHint>)args[0];
+ for (TaskWithLocationHint task : tasks) {
+ scheduledTasks.add(task.getTaskIndex());
+ }
return null;
}}).when(mockContext).scheduleVertexTasks(anyList());
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-tests/src/test/java/org/apache/tez/mapreduce/TestMRRJobsDAGApi.java
----------------------------------------------------------------------
diff --git a/tez-tests/src/test/java/org/apache/tez/mapreduce/TestMRRJobsDAGApi.java b/tez-tests/src/test/java/org/apache/tez/mapreduce/TestMRRJobsDAGApi.java
index 1fc362c..b690a11 100644
--- a/tez-tests/src/test/java/org/apache/tez/mapreduce/TestMRRJobsDAGApi.java
+++ b/tez-tests/src/test/java/org/apache/tez/mapreduce/TestMRRJobsDAGApi.java
@@ -88,6 +88,7 @@ import org.apache.tez.dag.api.client.DAGClient;
import org.apache.tez.dag.api.client.DAGStatus;
import org.apache.tez.dag.api.client.DAGStatus.State;
import org.apache.tez.mapreduce.common.MRInputAMSplitGenerator;
+import org.apache.tez.mapreduce.examples.BroadcastAndOneToOneExample;
import org.apache.tez.mapreduce.examples.MRRSleepJob;
import org.apache.tez.mapreduce.examples.MRRSleepJob.ISleepReducer;
import org.apache.tez.mapreduce.examples.MRRSleepJob.MRRSleepJobPartitioner;
@@ -674,8 +675,8 @@ public class TestMRRJobsDAGApi {
@Test(timeout = 60000)
public void testVertexGroups() throws Exception {
LOG.info("Running Group Test");
- Path inPath = new Path(TEST_ROOT_DIR, "in");
- Path outPath = new Path(TEST_ROOT_DIR, "out");
+ Path inPath = new Path(TEST_ROOT_DIR, "in-groups");
+ Path outPath = new Path(TEST_ROOT_DIR, "out-groups");
FSDataOutputStream out = remoteFs.create(inPath);
OutputStreamWriter writer = new OutputStreamWriter(out);
writer.write("abcd ");
@@ -692,6 +693,17 @@ public class TestMRRJobsDAGApi {
throw new TezUncheckedException("VertexGroups Test Failed");
}
}
+
+ @Test(timeout = 60000)
+ public void testBroadcastAndOneToOne() throws Exception {
+ LOG.info("Running BroadcastAndOneToOne Test");
+ BroadcastAndOneToOneExample job = new BroadcastAndOneToOneExample();
+ if (job.run(mrrTezCluster.getConfig())) {
+ LOG.info("Success BroadcastAndOneToOne Test");
+ } else {
+ throw new TezUncheckedException("BroadcastAndOneToOne Test Failed");
+ }
+ }
// This class should not be used by more than one test in a single run, since
// the path it writes to is not dynamic.
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/0a681b34/tez-tests/src/test/java/org/apache/tez/test/dag/MultiAttemptDAG.java
----------------------------------------------------------------------
diff --git a/tez-tests/src/test/java/org/apache/tez/test/dag/MultiAttemptDAG.java b/tez-tests/src/test/java/org/apache/tez/test/dag/MultiAttemptDAG.java
index c288f7e..d485c85 100644
--- a/tez-tests/src/test/java/org/apache/tez/test/dag/MultiAttemptDAG.java
+++ b/tez-tests/src/test/java/org/apache/tez/test/dag/MultiAttemptDAG.java
@@ -18,6 +18,7 @@
package org.apache.tez.test.dag;
+import com.google.common.collect.Lists;
import com.google.common.primitives.Ints;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
@@ -35,6 +36,7 @@ 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.Event;
import org.apache.tez.runtime.api.LogicalInput;
@@ -53,7 +55,6 @@ import org.apache.tez.test.TestInput;
import org.apache.tez.test.TestOutput;
import org.apache.tez.test.TestProcessor;
-import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
@@ -118,9 +119,9 @@ public class MultiAttemptDAG {
} else if (successAttemptId == context.getDAGAttemptNumber()) {
LOG.info("Scheduling tasks for vertex=" + context.getVertexName());
int numTasks = context.getVertexNumTasks(context.getVertexName());
- List<Integer> scheduledTasks = new ArrayList<Integer>(numTasks);
+ List<TaskWithLocationHint> scheduledTasks = Lists.newArrayListWithCapacity(numTasks);
for (int i=0; i<numTasks; ++i) {
- scheduledTasks.add(new Integer(i));
+ scheduledTasks.add(new TaskWithLocationHint(new Integer(i), null));
}
context.scheduleVertexTasks(scheduledTasks);
}