You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tez.apache.org by ss...@apache.org on 2014/04/16 00:22:36 UTC
git commit: TEZ-708. Add a LocalTaskScheduler for use in Local mode.
Contributed by Jonathan Eagles.
Repository: incubator-tez
Updated Branches:
refs/heads/master de87def6b -> 98fd541b3
TEZ-708. Add a LocalTaskScheduler for use in Local mode. Contributed by
Jonathan Eagles.
Project: http://git-wip-us.apache.org/repos/asf/incubator-tez/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-tez/commit/98fd541b
Tree: http://git-wip-us.apache.org/repos/asf/incubator-tez/tree/98fd541b
Diff: http://git-wip-us.apache.org/repos/asf/incubator-tez/diff/98fd541b
Branch: refs/heads/master
Commit: 98fd541b31dae3c5a3ec748c7d6ebe4329cd0769
Parents: de87def
Author: Siddharth Seth <ss...@apache.org>
Authored: Tue Apr 15 15:21:05 2014 -0700
Committer: Siddharth Seth <ss...@apache.org>
Committed: Tue Apr 15 15:21:05 2014 -0700
----------------------------------------------------------------------
.../tez/dag/app/rm/LocalTaskScheduler.java | 330 +++++++++++++++++++
.../apache/tez/dag/app/rm/TaskScheduler.java | 12 +-
.../dag/app/rm/TaskSchedulerEventHandler.java | 23 +-
.../tez/dag/app/rm/TaskSchedulerInterface.java | 49 +++
.../tez/dag/app/rm/TestLocalTaskScheduler.java | 95 ++++++
.../dag/app/rm/TestTaskSchedulerHelpers.java | 11 +-
6 files changed, 507 insertions(+), 13 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/98fd541b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/LocalTaskScheduler.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/LocalTaskScheduler.java b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/LocalTaskScheduler.java
new file mode 100644
index 0000000..7b81407
--- /dev/null
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/LocalTaskScheduler.java
@@ -0,0 +1,330 @@
+/**
+* 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.dag.app.rm;
+
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.PriorityBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.service.AbstractService;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+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.hadoop.yarn.api.records.Priority;
+import org.apache.hadoop.yarn.api.records.Resource;
+
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.app.AppContext;
+import org.apache.tez.dag.app.rm.TaskScheduler.TaskSchedulerAppCallback;
+import org.apache.tez.dag.app.rm.container.ContainerSignatureMatcher;
+
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+
+public class LocalTaskScheduler extends AbstractService implements TaskSchedulerInterface {
+
+ private static final Log LOG = LogFactory.getLog(LocalTaskScheduler.class);
+
+ final TaskSchedulerAppCallback realAppClient;
+ final TaskSchedulerAppCallback appClientDelegate;
+ final ContainerSignatureMatcher containerSignatureMatcher;
+ final PriorityBlockingQueue<TaskRequest> taskRequestQueue;
+ AsyncDelegateRequestHandler taskRequestHandler;
+ Thread asyncDelegateRequestThread;
+ final ExecutorService appCallbackExecutor;
+
+ final HashMap<Object, Container> taskAllocations;
+ final String appHostName;
+ final int appHostPort;
+ final String appTrackingUrl;
+ final AppContext appContext;
+
+ public LocalTaskScheduler(TaskSchedulerAppCallback appClient,
+ ContainerSignatureMatcher containerSignatureMatcher, String appHostName,
+ int appHostPort, String appTrackingUrl, AppContext appContext) {
+ super(LocalTaskScheduler.class.getName());
+ this.realAppClient = appClient;
+ this.appCallbackExecutor = createAppCallbackExecutorService();
+ this.containerSignatureMatcher = containerSignatureMatcher;
+ this.appClientDelegate = createAppCallbackDelegate(appClient);
+ this.appHostName = appHostName;
+ this.appHostPort = appHostPort;
+ this.appTrackingUrl = appTrackingUrl;
+ this.appContext = appContext;
+ taskRequestQueue = new PriorityBlockingQueue<TaskRequest>();
+ taskAllocations = new LinkedHashMap<Object, Container>();
+ }
+
+ private ExecutorService createAppCallbackExecutorService() {
+ return Executors.newSingleThreadExecutor(new ThreadFactoryBuilder()
+ .setNameFormat("TaskSchedulerAppCaller #%d").setDaemon(true).build());
+ }
+
+ private TaskSchedulerAppCallback createAppCallbackDelegate(
+ TaskSchedulerAppCallback realAppClient) {
+ return new TaskSchedulerAppCallbackWrapper(realAppClient,
+ appCallbackExecutor);
+ }
+
+ @Override
+ public Resource getAvailableResources() {
+ Resource freeResources = Resource.newInstance(
+ (int)Runtime.getRuntime().freeMemory()/(1024*1024),
+ Runtime.getRuntime().availableProcessors());
+ return freeResources;
+ }
+
+ @Override
+ public int getClusterNodeCount() {
+ return 1;
+ }
+
+ @Override
+ public void resetMatchLocalityForAllHeldContainers() {
+ }
+
+ @Override
+ public Resource getTotalResources() {
+ Resource totalResources = Resource.newInstance(
+ (int)Runtime.getRuntime().maxMemory()/(1024*1024),
+ Runtime.getRuntime().availableProcessors());
+ return totalResources;
+ }
+
+ @Override
+ public void blacklistNode(NodeId nodeId) {
+ }
+
+ @Override
+ public void unblacklistNode(NodeId nodeId) {
+ }
+
+ @Override
+ public void allocateTask(Object task, Resource capability, String[] hosts,
+ String[] racks, Priority priority, Object containerSignature,
+ Object clientCookie) {
+
+ taskRequestHandler.addAllocateTaskRequest(task, capability, priority, clientCookie);
+ }
+
+ @Override
+ public boolean deallocateTask(Object task, boolean taskSucceeded) {
+ return taskRequestHandler.addDeallocateTaskRequest(task);
+ }
+
+ @Override
+ public Object deallocateContainer(ContainerId containerId) {
+ return null;
+ }
+
+ @Override
+ public void serviceInit(Configuration conf) {
+ taskRequestHandler = new AsyncDelegateRequestHandler(taskRequestQueue,
+ new LocalContainerFactory(appContext),
+ taskAllocations,
+ appClientDelegate,
+ conf);
+ asyncDelegateRequestThread = new Thread(taskRequestHandler);
+ }
+
+ @Override
+ public void serviceStart() {
+ asyncDelegateRequestThread.start();
+ }
+
+ @Override
+ public void serviceStop() throws InterruptedException {
+ if (asyncDelegateRequestThread != null) {
+ asyncDelegateRequestThread.interrupt();
+ }
+ appCallbackExecutor.shutdownNow();
+ appCallbackExecutor.awaitTermination(1000l, TimeUnit.MILLISECONDS);
+ }
+
+ @Override
+ public void setShouldUnregister() {
+ }
+
+ static class LocalContainerFactory {
+ final AppContext appContext;
+ AtomicInteger nextId;
+
+ public LocalContainerFactory(AppContext appContext) {
+ this.appContext = appContext;
+ this.nextId = new AtomicInteger(1);
+ }
+
+ public Container createContainer(Resource capability, Priority priority) {
+ ApplicationAttemptId appAttemptId = appContext.getApplicationAttemptId();
+ ContainerId containerId = ContainerId.newInstance(appAttemptId, nextId.getAndIncrement());
+ NodeId nodeId = NodeId.newInstance("127.0.0.1", 0);
+ String nodeHttpAddress = "127.0.0.1:0";
+
+ Container container = Container.newInstance(containerId,
+ nodeId,
+ nodeHttpAddress,
+ capability,
+ priority,
+ null);
+
+ return container;
+ }
+ }
+
+ static class TaskRequest implements Comparable<TaskRequest> {
+ // Higher prority than Priority.UNDEFINED
+ static final int HIGHEST_PRIORITY = -2;
+ Object task;
+ Priority priority;
+
+ public TaskRequest(Object task, Priority priority) {
+ this.task = task;
+ this.priority = priority;
+ }
+
+ @Override
+ public int compareTo(TaskRequest request) {
+ return request.priority.compareTo(this.priority);
+ }
+ }
+
+ static class AllocateTaskRequest extends TaskRequest {
+ Resource capability;
+ Object clientCookie;
+
+ public AllocateTaskRequest(Object task, Resource capability, Priority priority,
+ Object clientCookie) {
+ super(task, priority);
+ this.capability = capability;
+ this.clientCookie = clientCookie;
+ }
+ }
+
+ static class DeallocateTaskRequest extends TaskRequest {
+ static final Priority DEALLOCATE_PRIORITY = Priority.newInstance(HIGHEST_PRIORITY);
+
+ public DeallocateTaskRequest(Object task) {
+ super(task, DEALLOCATE_PRIORITY);
+ }
+ }
+
+ static class AsyncDelegateRequestHandler implements Runnable {
+ final BlockingQueue<TaskRequest> taskRequestQueue;
+ final LocalContainerFactory localContainerFactory;
+ final HashMap<Object, Container> taskAllocations;
+ final TaskSchedulerAppCallback appClientDelegate;
+ final int MAX_TASKS;
+
+ AsyncDelegateRequestHandler(BlockingQueue<TaskRequest> taskRequestQueue,
+ LocalContainerFactory localContainerFactory,
+ HashMap<Object, Container> taskAllocations,
+ TaskSchedulerAppCallback appClientDelegate,
+ Configuration conf) {
+ this.taskRequestQueue = taskRequestQueue;
+ this.localContainerFactory = localContainerFactory;
+ this.taskAllocations = taskAllocations;
+ this.appClientDelegate = appClientDelegate;
+ this.MAX_TASKS = conf.getInt(TezConfiguration.TEZ_AM_INLINE_TASK_EXECUTION_MAX_TASKS,
+ TezConfiguration.TEZ_AM_INLINE_TASK_EXECUTION_MAX_TASKS_DEFAULT);
+ }
+
+ public void addAllocateTaskRequest(Object task, Resource capability, Priority priority,
+ Object clientCookie) {
+ try {
+ taskRequestQueue.put(new AllocateTaskRequest(task, capability, priority, clientCookie));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ }
+
+ public boolean addDeallocateTaskRequest(Object task) {
+ try {
+ taskRequestQueue.put(new DeallocateTaskRequest(task));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ synchronized(taskRequestQueue) {
+ taskRequestQueue.notify();
+ }
+ return true;
+ }
+
+ boolean shouldWait() {
+ return taskAllocations.size() >= MAX_TASKS;
+ }
+
+ @Override
+ public void run() {
+ while(true) {
+ synchronized(taskRequestQueue) {
+ try {
+ if (shouldWait()) {
+ taskRequestQueue.wait();
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ processRequest();
+ }
+ }
+
+ void processRequest() {
+ try {
+ TaskRequest request = taskRequestQueue.take();
+ if (request instanceof AllocateTaskRequest) {
+ allocateTask((AllocateTaskRequest)request);
+ }
+ else if (request instanceof DeallocateTaskRequest) {
+ deallocateTask((DeallocateTaskRequest)request);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ } catch (NullPointerException e) {
+ LOG.warn("Task request was badly constructed");
+ }
+ }
+
+ void allocateTask(AllocateTaskRequest request) {
+ Container container = localContainerFactory.createContainer(request.capability,
+ request.priority);
+ taskAllocations.put(request.task, container);
+ appClientDelegate.taskAllocated(request.task, request.clientCookie, container);
+ }
+
+ void deallocateTask(DeallocateTaskRequest request) {
+ Container container = taskAllocations.remove(request.task);
+ if (container != null) {
+ appClientDelegate.containerBeingReleased(container.getId());
+ }
+ else {
+ LOG.warn("Unable to find and remove task " + request.task + " from task allocations");
+ }
+ }
+ }
+}
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/98fd541b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskScheduler.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskScheduler.java b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskScheduler.java
index 48f67b9..13a71d8 100644
--- a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskScheduler.java
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskScheduler.java
@@ -80,7 +80,7 @@ import com.google.common.util.concurrent.ThreadFactoryBuilder;
}
*/
public class TaskScheduler extends AbstractService
- implements AMRMClientAsync.CallbackHandler {
+ implements AMRMClientAsync.CallbackHandler, TaskSchedulerInterface {
private static final Log LOG = LogFactory.getLog(TaskScheduler.class);
public interface TaskSchedulerAppCallback {
@@ -260,10 +260,12 @@ public class TaskScheduler extends AbstractService
.setNameFormat("TaskSchedulerAppCaller #%d").setDaemon(true).build());
}
+ @Override
public Resource getAvailableResources() {
return amRmClient.getAvailableResources();
}
+ @Override
public int getClusterNodeCount() {
// this can potentially be cheaper after YARN-1722
return amRmClient.getClusterNodeCount();
@@ -275,6 +277,7 @@ public class TaskScheduler extends AbstractService
appCallbackExecutor);
}
+ @Override
public void setShouldUnregister() {
this.shouldUnregister.set(true);
}
@@ -738,6 +741,7 @@ public class TaskScheduler extends AbstractService
return null;
}
+ @Override
public synchronized void resetMatchLocalityForAllHeldContainers() {
for (HeldContainer heldContainer : heldContainers.values()) {
heldContainer.resetLocalityMatchLevel();
@@ -795,16 +799,19 @@ public class TaskScheduler extends AbstractService
appClientDelegate.onError(t);
}
+ @Override
public Resource getTotalResources() {
return totalResources;
}
+ @Override
public synchronized void blacklistNode(NodeId nodeId) {
LOG.info("Blacklisting node: " + nodeId);
amRmClient.addNodeToBlacklist(nodeId);
blacklistedNodes.add(nodeId);
}
+ @Override
public synchronized void unblacklistNode(NodeId nodeId) {
if (blacklistedNodes.remove(nodeId)) {
LOG.info("UnBlacklisting node: " + nodeId);
@@ -812,6 +819,7 @@ public class TaskScheduler extends AbstractService
}
}
+ @Override
public synchronized void allocateTask(
Object task,
Resource capability,
@@ -844,6 +852,7 @@ public class TaskScheduler extends AbstractService
* specify whether the task succeeded or failed.
* @return true if a container is assigned to this task.
*/
+ @Override
public boolean deallocateTask(Object task, boolean taskSucceeded) {
Map<CookieContainerRequest, Container> assignedContainers = null;
@@ -897,6 +906,7 @@ public class TaskScheduler extends AbstractService
return true;
}
+ @Override
public synchronized Object deallocateContainer(ContainerId containerId) {
Object task = unAssignContainer(containerId, true);
if(task != null) {
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/98fd541b/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 108cbdf..8fdcfc2 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
@@ -39,6 +39,7 @@ import org.apache.hadoop.yarn.api.records.NodeReport;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.event.Event;
import org.apache.hadoop.yarn.event.EventHandler;
+import org.apache.tez.dag.api.TezConfiguration;
import org.apache.tez.dag.api.TezUncheckedException;
import org.apache.tez.dag.api.VertexLocationHint.TaskLocationHint;
import org.apache.tez.dag.api.client.DAGClientServer;
@@ -74,7 +75,7 @@ public class TaskSchedulerEventHandler extends AbstractService
protected final AppContext appContext;
@SuppressWarnings("rawtypes")
private final EventHandler eventHandler;
- protected TaskScheduler taskScheduler;
+ protected TaskSchedulerInterface taskScheduler;
private DAGAppMaster dagAppMaster;
private Map<ApplicationAccessType, String> appAcls = null;
private Thread eventHandlingThread;
@@ -347,10 +348,18 @@ public class TaskSchedulerEventHandler extends AbstractService
}
- protected TaskScheduler createTaskScheduler(String host, int port,
+ protected TaskSchedulerInterface createTaskScheduler(String host, int port,
String trackingUrl, AppContext appContext) {
- return new TaskScheduler(this, this.containerSignatureMatcher,
- host, port, trackingUrl, appContext);
+ boolean isLocal = getConfig().getBoolean(TezConfiguration.TEZ_LOCAL_MODE,
+ TezConfiguration.TEZ_LOCAL_MODE_DEFAULT);
+ if (isLocal) {
+ return new LocalTaskScheduler(this, this.containerSignatureMatcher,
+ host, port, trackingUrl, appContext);
+ }
+ else {
+ return new TaskScheduler(this, this.containerSignatureMatcher,
+ host, port, trackingUrl, appContext);
+ }
}
@Override
@@ -359,8 +368,8 @@ public class TaskSchedulerEventHandler extends AbstractService
dagAppMaster = appContext.getAppMaster();
taskScheduler = createTaskScheduler(serviceAddr.getHostName(),
serviceAddr.getPort(), "", appContext);
- taskScheduler.init(getConfig());
- taskScheduler.start();
+ ((AbstractService)taskScheduler).init(getConfig());
+ ((AbstractService)taskScheduler).start();
if (shouldUnregisterFlag.get()) {
// Flag may have been set earlier when task scheduler was not initialized
taskScheduler.setShouldUnregister();
@@ -405,7 +414,7 @@ public class TaskSchedulerEventHandler extends AbstractService
eventHandlingThread.interrupt();
}
if (taskScheduler != null) {
- taskScheduler.stop();
+ ((AbstractService)taskScheduler).stop();
}
}
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/98fd541b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerInterface.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerInterface.java b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerInterface.java
new file mode 100644
index 0000000..5b6f179
--- /dev/null
+++ b/tez-dag/src/main/java/org/apache/tez/dag/app/rm/TaskSchedulerInterface.java
@@ -0,0 +1,49 @@
+/**
+* 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.dag.app.rm;
+
+import org.apache.hadoop.yarn.api.records.ContainerId;
+import org.apache.hadoop.yarn.api.records.NodeId;
+import org.apache.hadoop.yarn.api.records.Priority;
+import org.apache.hadoop.yarn.api.records.Resource;
+
+public interface TaskSchedulerInterface {
+
+ public abstract Resource getAvailableResources();
+
+ public abstract int getClusterNodeCount();
+
+ public abstract void resetMatchLocalityForAllHeldContainers();
+
+ public abstract Resource getTotalResources();
+
+ public abstract void blacklistNode(NodeId nodeId);
+
+ public abstract void unblacklistNode(NodeId nodeId);
+
+ public abstract void allocateTask(Object task, Resource capability,
+ String[] hosts, String[] racks, Priority priority,
+ Object containerSignature, Object clientCookie);
+
+ public abstract boolean deallocateTask(Object task, boolean taskSucceeded);
+
+ public abstract Object deallocateContainer(ContainerId containerId);
+
+ public abstract void setShouldUnregister();
+}
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/98fd541b/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestLocalTaskScheduler.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestLocalTaskScheduler.java b/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestLocalTaskScheduler.java
new file mode 100644
index 0000000..1fb393f
--- /dev/null
+++ b/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestLocalTaskScheduler.java
@@ -0,0 +1,95 @@
+/**
+ * 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.dag.app.rm;
+
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.concurrent.PriorityBlockingQueue;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+import org.apache.hadoop.yarn.api.records.ApplicationId;
+import org.apache.hadoop.yarn.api.records.Container;
+import org.apache.hadoop.yarn.api.records.Priority;
+
+import org.apache.tez.dag.api.TezConfiguration;
+import org.apache.tez.dag.app.AppContext;
+import org.apache.tez.dag.app.rm.LocalTaskScheduler.AsyncDelegateRequestHandler;
+import org.apache.tez.dag.app.rm.LocalTaskScheduler.LocalContainerFactory;
+import org.apache.tez.dag.app.rm.LocalTaskScheduler.TaskRequest;
+import org.apache.tez.dag.app.rm.TaskScheduler.TaskSchedulerAppCallback;
+
+public class TestLocalTaskScheduler {
+
+ public AppContext createMockAppContext() {
+
+ ApplicationId appId = ApplicationId.newInstance(2000, 1);
+ ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, 1);
+
+ AppContext appContext = mock(AppContext.class);
+ doReturn(appAttemptId).when(appContext).getApplicationAttemptId();
+
+ return appContext;
+ }
+
+ @Test
+ public void maxTasksAllocationsCannotBeExceeded() {
+
+ final int MAX_TASKS = 4;
+ TezConfiguration tezConf = new TezConfiguration();
+ tezConf.setInt(TezConfiguration.TEZ_AM_INLINE_TASK_EXECUTION_MAX_TASKS, MAX_TASKS);
+
+ LocalContainerFactory containerFactory = new LocalContainerFactory(createMockAppContext());
+ HashMap<Object, Container> taskAllocations = new LinkedHashMap<Object, Container>();
+ PriorityBlockingQueue<TaskRequest> taskRequestQueue = new PriorityBlockingQueue<TaskRequest>();
+ TaskSchedulerAppCallback appClientDelegate = mock(TaskSchedulerAppCallback.class);
+
+ // Object under test
+ AsyncDelegateRequestHandler requestHandler =
+ new AsyncDelegateRequestHandler(taskRequestQueue,
+ containerFactory,
+ taskAllocations,
+ appClientDelegate,
+ tezConf);
+
+ // Allocate up to max tasks
+ for (int i = 0; i < MAX_TASKS; i++) {
+ Priority priority = Priority.newInstance(20);
+ requestHandler.addAllocateTaskRequest(new Long(i), null, priority, null);
+ requestHandler.processRequest();
+ }
+
+ // Only MAX_TASKS number of tasks should have been allocated
+ Assert.assertEquals("Wrong number of allocate tasks", MAX_TASKS, taskAllocations.size());
+ Assert.assertTrue("Another allocation should not fit", requestHandler.shouldWait());
+
+ // Deallocate down to zero
+ for (int i = 0; i < MAX_TASKS; i++) {
+ requestHandler.addDeallocateTaskRequest(new Long(i));
+ requestHandler.processRequest();
+ }
+
+ // All allocated tasks should have been removed
+ Assert.assertEquals("Wrong number of allocate tasks", 0, taskAllocations.size());
+ }
+}
http://git-wip-us.apache.org/repos/asf/incubator-tez/blob/98fd541b/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerHelpers.java
----------------------------------------------------------------------
diff --git a/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerHelpers.java b/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerHelpers.java
index 72cc8e7..a9fdaab 100644
--- a/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerHelpers.java
+++ b/tez-dag/src/test/java/org/apache/tez/dag/app/rm/TestTaskSchedulerHelpers.java
@@ -40,6 +40,7 @@ import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
+import org.apache.hadoop.service.AbstractService;
import org.apache.hadoop.yarn.api.protocolrecords.RegisterApplicationMasterResponse;
import org.apache.hadoop.yarn.api.records.ApplicationAccessType;
import org.apache.hadoop.yarn.api.records.Container;
@@ -132,24 +133,24 @@ class TestTaskSchedulerHelpers {
}
@Override
- public TaskScheduler createTaskScheduler(String host, int port,
+ public TaskSchedulerInterface createTaskScheduler(String host, int port,
String trackingUrl, AppContext appContext) {
return new TaskSchedulerWithDrainableAppCallback(this,
containerSignatureMatcher, host, port, trackingUrl, amrmClientAsync,
appContext);
}
- public TaskScheduler getSpyTaskScheduler() {
+ public TaskSchedulerInterface getSpyTaskScheduler() {
return this.taskScheduler;
}
@Override
public void serviceStart() {
- TaskScheduler taskSchedulerReal = createTaskScheduler("host", 0, "",
+ TaskSchedulerInterface taskSchedulerReal = createTaskScheduler("host", 0, "",
appContext);
// Init the service so that reuse configuration is picked up.
- taskSchedulerReal.serviceInit(getConfig());
- taskSchedulerReal.serviceStart();
+ ((AbstractService)taskSchedulerReal).init(getConfig());
+ ((AbstractService)taskSchedulerReal).start();
taskScheduler = spy(taskSchedulerReal);
}