use of org.apache.tez.dag.api.records.DAGProtos.PlanLocalResource in project tez by apache.
the class TestVertexImpl method setupPostDagCreation.
@SuppressWarnings({ "unchecked", "rawtypes" })
public void setupPostDagCreation() throws TezException {
String dagName = "dag0";
// dispatcher may be created multiple times (setupPostDagCreation may be called multiples)
if (dispatcher != null) {
dispatcher.stop();
}
dispatcher = new DrainDispatcher();
appContext = mock(AppContext.class);
when(appContext.getHadoopShim()).thenReturn(new DefaultHadoopShim());
when(appContext.getContainerLauncherName(anyInt())).thenReturn(TezConstants.getTezYarnServicePluginName());
thh = mock(TaskHeartbeatHandler.class);
historyEventHandler = mock(HistoryEventHandler.class);
TaskSchedulerManager taskScheduler = mock(TaskSchedulerManager.class);
UserGroupInformation ugi;
try {
ugi = UserGroupInformation.getCurrentUser();
} catch (IOException e) {
throw new RuntimeException(e);
}
DAG dag = mock(DAG.class);
doReturn(ugi).when(dag).getDagUGI();
doReturn(dagName).when(dag).getName();
Map<String, LocalResource> localResources = new HashMap<>();
for (PlanLocalResource planLR : dagPlan.getLocalResourceList()) {
localResources.put(planLR.getName(), DagTypeConverters.convertPlanLocalResourceToLocalResource(planLR));
}
when(dag.getLocalResources()).thenReturn(localResources);
doReturn(appAttemptId).when(appContext).getApplicationAttemptId();
doReturn(appAttemptId.getApplicationId()).when(appContext).getApplicationID();
doReturn(dag).when(appContext).getCurrentDAG();
execService = mock(ListeningExecutorService.class);
final ListenableFuture<Void> mockFuture = mock(ListenableFuture.class);
Mockito.doAnswer(new Answer() {
public ListenableFuture<Void> answer(InvocationOnMock invocation) {
Object[] args = invocation.getArguments();
CallableEvent e = (CallableEvent) args[0];
dispatcher.getEventHandler().handle(e);
return mockFuture;
}
}).when(execService).submit((Callable<Void>) any());
MockClock clock = new MockClock();
doReturn(execService).when(appContext).getExecService();
doReturn(conf).when(appContext).getAMConf();
doReturn(new Credentials()).when(dag).getCredentials();
doReturn(DAGPlan.getDefaultInstance()).when(dag).getJobPlan();
doReturn(dagId).when(appContext).getCurrentDAGID();
doReturn(dagId).when(dag).getID();
doReturn(taskScheduler).when(appContext).getTaskScheduler();
doReturn(Resource.newInstance(102400, 60)).when(taskScheduler).getTotalResources(0);
doReturn(historyEventHandler).when(appContext).getHistoryHandler();
doReturn(dispatcher.getEventHandler()).when(appContext).getEventHandler();
doReturn(clock).when(appContext).getClock();
vertexGroups = Maps.newHashMap();
for (PlanVertexGroupInfo groupInfo : dagPlan.getVertexGroupsList()) {
vertexGroups.put(groupInfo.getGroupName(), new VertexGroupInfo(groupInfo));
}
// updateTracker may be created multiple times (setupPostDagCreation may be called multiples)
if (updateTracker != null) {
updateTracker.stop();
}
updateTracker = new StateChangeNotifierForTest(appContext.getCurrentDAG());
setupVertices();
when(dag.getVertex(any(TezVertexID.class))).thenAnswer(new Answer<Vertex>() {
@Override
public Vertex answer(InvocationOnMock invocation) throws Throwable {
Object[] args = invocation.getArguments();
if (args.length != 1) {
return null;
}
TezVertexID vId = (TezVertexID) args[0];
return vertexIdMap.get(vId);
}
});
when(dag.getVertex(any(String.class))).thenAnswer(new Answer<Vertex>() {
@Override
public Vertex answer(InvocationOnMock invocation) throws Throwable {
Object[] args = invocation.getArguments();
if (args.length != 1) {
return null;
}
String vId = (String) args[0];
return vertices.get(vId);
}
});
// TODO - this test logic is tightly linked to impl DAGImpl code.
edges = new HashMap<String, Edge>();
for (EdgePlan edgePlan : dagPlan.getEdgeList()) {
EdgeProperty edgeProperty = DagTypeConverters.createEdgePropertyMapFromDAGPlan(edgePlan);
edges.put(edgePlan.getId(), new Edge(edgeProperty, dispatcher.getEventHandler(), conf));
}
parseVertexEdges();
for (Edge edge : edges.values()) {
edge.initialize();
}
dispatcher.register(CallableEventType.class, new CallableEventDispatcher());
taskAttemptEventDispatcher = new TaskAttemptEventDispatcher();
dispatcher.register(TaskAttemptEventType.class, taskAttemptEventDispatcher);
taskEventDispatcher = new TaskEventDispatcher();
dispatcher.register(TaskEventType.class, taskEventDispatcher);
vertexEventDispatcher = new VertexEventDispatcher();
dispatcher.register(VertexEventType.class, vertexEventDispatcher);
dagEventDispatcher = new DagEventDispatcher();
dispatcher.register(DAGEventType.class, dagEventDispatcher);
amSchedulerEventDispatcher = new AMSchedulerEventDispatcher();
dispatcher.register(AMSchedulerEventType.class, amSchedulerEventDispatcher);
dispatcher.init(conf);
dispatcher.start();
}
use of org.apache.tez.dag.api.records.DAGProtos.PlanLocalResource in project tez by apache.
the class DagTypeConverters method convertFromPlanLocalResources.
public static Map<String, LocalResource> convertFromPlanLocalResources(PlanLocalResourcesProto proto) {
Map<String, LocalResource> localResources = new HashMap<String, LocalResource>(proto.getLocalResourcesCount());
for (PlanLocalResource plr : proto.getLocalResourcesList()) {
String name = plr.getName();
LocalResource lr = convertPlanLocalResourceToLocalResource(plr);
localResources.put(name, lr);
}
return localResources;
}
use of org.apache.tez.dag.api.records.DAGProtos.PlanLocalResource in project tez by apache.
the class DagTypeConverters method createLocalResourceMapFromDAGPlan.
public static Map<String, LocalResource> createLocalResourceMapFromDAGPlan(List<PlanLocalResource> localResourcesList) {
Map<String, LocalResource> map = new HashMap<String, LocalResource>();
for (PlanLocalResource res : localResourcesList) {
LocalResource r = new LocalResourcePBImpl();
// else we will receive a default value back, eg ""
if (res.hasPattern()) {
r.setPattern(res.getPattern());
}
r.setResource(convertToYarnURL(res.getUri()));
r.setSize(res.getSize());
r.setTimestamp(res.getTimeStamp());
r.setType(DagTypeConverters.convertFromDAGPlan(res.getType()));
r.setVisibility(DagTypeConverters.convertFromDAGPlan(res.getVisibility()));
map.put(res.getName(), r);
}
return map;
}
use of org.apache.tez.dag.api.records.DAGProtos.PlanLocalResource in project tez by apache.
the class DagTypeConverters method convertToDAGPlan.
public static List<PlanLocalResource> convertToDAGPlan(Map<String, LocalResource> lrs) {
List<PlanLocalResource> planLrs = Lists.newArrayListWithCapacity(lrs.size());
for (Entry<String, LocalResource> entry : lrs.entrySet()) {
PlanLocalResource.Builder localResourcesBuilder = PlanLocalResource.newBuilder();
String key = entry.getKey();
LocalResource lr = entry.getValue();
localResourcesBuilder.setName(key);
localResourcesBuilder.setUri(DagTypeConverters.convertToDAGPlan(lr.getResource()));
localResourcesBuilder.setSize(lr.getSize());
localResourcesBuilder.setTimeStamp(lr.getTimestamp());
localResourcesBuilder.setType(DagTypeConverters.convertToDAGPlan(lr.getType()));
localResourcesBuilder.setVisibility(DagTypeConverters.convertToDAGPlan(lr.getVisibility()));
if (lr.getType() == LocalResourceType.PATTERN) {
if (lr.getPattern() == null || lr.getPattern().isEmpty()) {
throw new TezUncheckedException("LocalResource type set to pattern" + " but pattern is null or empty");
}
localResourcesBuilder.setPattern(lr.getPattern());
}
planLrs.add(localResourcesBuilder.build());
}
return planLrs;
}
use of org.apache.tez.dag.api.records.DAGProtos.PlanLocalResource in project tez by apache.
the class DagTypeConverters method convertFromLocalResources.
public static PlanLocalResourcesProto convertFromLocalResources(Map<String, LocalResource> localResources) {
PlanLocalResourcesProto.Builder builder = PlanLocalResourcesProto.newBuilder();
for (Map.Entry<String, LocalResource> entry : localResources.entrySet()) {
PlanLocalResource plr = convertLocalResourceToPlanLocalResource(entry.getKey(), entry.getValue());
builder.addLocalResources(plr);
}
return builder.build();
}
Aggregations