You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by pw...@apache.org on 2014/01/04 01:31:31 UTC

[4/6] merge upstream/master

http://git-wip-us.apache.org/repos/asf/incubator-spark/blob/8ddbd531/yarn/alpha/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocationHandler.scala
----------------------------------------------------------------------
diff --cc yarn/alpha/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocationHandler.scala
index 0000000,c8af653..e91257b
mode 000000,100644..100644
--- a/yarn/alpha/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocationHandler.scala
+++ b/yarn/alpha/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocationHandler.scala
@@@ -1,0 -1,680 +1,680 @@@
+ /*
+  * 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.spark.deploy.yarn
+ 
+ import java.lang.{Boolean => JBoolean}
+ import java.util.{Collections, Set => JSet}
+ import java.util.concurrent.{CopyOnWriteArrayList, ConcurrentHashMap}
+ import java.util.concurrent.atomic.AtomicInteger
+ 
+ import scala.collection
+ import scala.collection.JavaConversions._
+ import scala.collection.mutable.{ArrayBuffer, HashMap, HashSet}
+ 
+ import org.apache.spark.{Logging, SparkConf}
+ import org.apache.spark.scheduler.{SplitInfo,TaskSchedulerImpl}
+ import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend
+ import org.apache.spark.util.Utils
+ 
+ import org.apache.hadoop.conf.Configuration
+ import org.apache.hadoop.yarn.api.AMRMProtocol
+ import org.apache.hadoop.yarn.api.records.{AMResponse, ApplicationAttemptId}
+ import org.apache.hadoop.yarn.api.records.{Container, ContainerId, ContainerStatus}
+ import org.apache.hadoop.yarn.api.records.{Priority, Resource, ResourceRequest}
+ import org.apache.hadoop.yarn.api.protocolrecords.{AllocateRequest, AllocateResponse}
+ import org.apache.hadoop.yarn.util.{RackResolver, Records}
+ 
+ 
+ object AllocationType extends Enumeration {
+   type AllocationType = Value
+   val HOST, RACK, ANY = Value
+ }
+ 
+ // TODO:
+ // Too many params.
+ // Needs to be mt-safe
+ // Need to refactor this to make it 'cleaner' ... right now, all computation is reactive - should
+ // make it more proactive and decoupled.
+ 
+ // Note that right now, we assume all node asks as uniform in terms of capabilities and priority
+ // Refer to http://developer.yahoo.com/blogs/hadoop/posts/2011/03/mapreduce-nextgen-scheduler/ for
+ // more info on how we are requesting for containers.
+ private[yarn] class YarnAllocationHandler(
+     val conf: Configuration,
+     val resourceManager: AMRMProtocol, 
+     val appAttemptId: ApplicationAttemptId,
+     val maxWorkers: Int,
+     val workerMemory: Int,
+     val workerCores: Int,
+     val preferredHostToCount: Map[String, Int], 
+     val preferredRackToCount: Map[String, Int],
+     val sparkConf: SparkConf)
+   extends Logging {
+   // These three are locked on allocatedHostToContainersMap. Complementary data structures
+   // allocatedHostToContainersMap : containers which are running : host, Set<containerid>
+   // allocatedContainerToHostMap: container to host mapping.
+   private val allocatedHostToContainersMap =
+     new HashMap[String, collection.mutable.Set[ContainerId]]()
+ 
+   private val allocatedContainerToHostMap = new HashMap[ContainerId, String]()
+ 
+   // allocatedRackCount is populated ONLY if allocation happens (or decremented if this is an
+   // allocated node)
+   // As with the two data structures above, tightly coupled with them, and to be locked on
+   // allocatedHostToContainersMap
+   private val allocatedRackCount = new HashMap[String, Int]()
+ 
+   // Containers which have been released.
+   private val releasedContainerList = new CopyOnWriteArrayList[ContainerId]()
+   // Containers to be released in next request to RM
+   private val pendingReleaseContainers = new ConcurrentHashMap[ContainerId, Boolean]
+ 
+   private val numWorkersRunning = new AtomicInteger()
+   // Used to generate a unique id per worker
+   private val workerIdCounter = new AtomicInteger()
+   private val lastResponseId = new AtomicInteger()
+   private val numWorkersFailed = new AtomicInteger()
+ 
+   def getNumWorkersRunning: Int = numWorkersRunning.intValue
+ 
+   def getNumWorkersFailed: Int = numWorkersFailed.intValue
+ 
+   def isResourceConstraintSatisfied(container: Container): Boolean = {
+     container.getResource.getMemory >= (workerMemory + YarnAllocationHandler.MEMORY_OVERHEAD)
+   }
+ 
+   def allocateContainers(workersToRequest: Int) {
+     // We need to send the request only once from what I understand ... but for now, not modifying
+     // this much.
+ 
+     // Keep polling the Resource Manager for containers
+     val amResp = allocateWorkerResources(workersToRequest).getAMResponse
+ 
+     val _allocatedContainers = amResp.getAllocatedContainers()
+ 
+     if (_allocatedContainers.size > 0) {
+       logDebug("""
+         Allocated containers: %d
+         Current worker count: %d
+         Containers released: %s
+         Containers to be released: %s
+         Cluster resources: %s
+         """.format(
+           _allocatedContainers.size,
+           numWorkersRunning.get(),
+           releasedContainerList,
+           pendingReleaseContainers,
+           amResp.getAvailableResources))
+ 
+       val hostToContainers = new HashMap[String, ArrayBuffer[Container]]()
+ 
+       // Ignore if not satisfying constraints      {
+       for (container <- _allocatedContainers) {
+         if (isResourceConstraintSatisfied(container)) {
+           // allocatedContainers += container
+ 
+           val host = container.getNodeId.getHost
+           val containers = hostToContainers.getOrElseUpdate(host, new ArrayBuffer[Container]())
+ 
+           containers += container
+         }
+         // Add all ignored containers to released list
+         else releasedContainerList.add(container.getId())
+       }
+ 
+       // Find the appropriate containers to use. Slightly non trivial groupBy ...
+       val dataLocalContainers = new HashMap[String, ArrayBuffer[Container]]()
+       val rackLocalContainers = new HashMap[String, ArrayBuffer[Container]]()
+       val offRackContainers = new HashMap[String, ArrayBuffer[Container]]()
+ 
+       for (candidateHost <- hostToContainers.keySet)
+       {
+         val maxExpectedHostCount = preferredHostToCount.getOrElse(candidateHost, 0)
+         val requiredHostCount = maxExpectedHostCount - allocatedContainersOnHost(candidateHost)
+ 
+         var remainingContainers = hostToContainers.get(candidateHost).getOrElse(null)
+         assert(remainingContainers != null)
+ 
+         if (requiredHostCount >= remainingContainers.size){
+           // Since we got <= required containers, add all to dataLocalContainers
+           dataLocalContainers.put(candidateHost, remainingContainers)
+           // all consumed
+           remainingContainers = null
+         }
+         else if (requiredHostCount > 0) {
+           // Container list has more containers than we need for data locality.
+           // Split into two : data local container count of (remainingContainers.size -
+           // requiredHostCount) and rest as remainingContainer
+           val (dataLocal, remaining) = remainingContainers.splitAt(
+             remainingContainers.size - requiredHostCount)
+           dataLocalContainers.put(candidateHost, dataLocal)
+           // remainingContainers = remaining
+ 
+           // yarn has nasty habit of allocating a tonne of containers on a host - discourage this :
+           // add remaining to release list. If we have insufficient containers, next allocation 
+           // cycle will reallocate (but wont treat it as data local)
+           for (container <- remaining) releasedContainerList.add(container.getId())
+           remainingContainers = null
+         }
+ 
+         // Now rack local
+         if (remainingContainers != null){
+           val rack = YarnAllocationHandler.lookupRack(conf, candidateHost)
+ 
+           if (rack != null){
+             val maxExpectedRackCount = preferredRackToCount.getOrElse(rack, 0)
+             val requiredRackCount = maxExpectedRackCount - allocatedContainersOnRack(rack) - 
+               rackLocalContainers.get(rack).getOrElse(List()).size
+ 
+ 
+             if (requiredRackCount >= remainingContainers.size){
+               // Add all to dataLocalContainers
+               dataLocalContainers.put(rack, remainingContainers)
+               // All consumed
+               remainingContainers = null
+             }
+             else if (requiredRackCount > 0) {
+               // container list has more containers than we need for data locality.
+               // Split into two : data local container count of (remainingContainers.size -
+               // requiredRackCount) and rest as remainingContainer
+               val (rackLocal, remaining) = remainingContainers.splitAt(
+                 remainingContainers.size - requiredRackCount)
+               val existingRackLocal = rackLocalContainers.getOrElseUpdate(rack,
+                 new ArrayBuffer[Container]())
+ 
+               existingRackLocal ++= rackLocal
+               remainingContainers = remaining
+             }
+           }
+         }
+ 
+         // If still not consumed, then it is off rack host - add to that list.
+         if (remainingContainers != null){
+           offRackContainers.put(candidateHost, remainingContainers)
+         }
+       }
+ 
+       // Now that we have split the containers into various groups, go through them in order : 
+       // first host local, then rack local and then off rack (everything else).
+       // Note that the list we create below tries to ensure that not all containers end up within a
+       // host if there are sufficiently large number of hosts/containers.
+ 
+       val allocatedContainers = new ArrayBuffer[Container](_allocatedContainers.size)
+       allocatedContainers ++= TaskSchedulerImpl.prioritizeContainers(dataLocalContainers)
+       allocatedContainers ++= TaskSchedulerImpl.prioritizeContainers(rackLocalContainers)
+       allocatedContainers ++= TaskSchedulerImpl.prioritizeContainers(offRackContainers)
+ 
+       // Run each of the allocated containers
+       for (container <- allocatedContainers) {
+         val numWorkersRunningNow = numWorkersRunning.incrementAndGet()
+         val workerHostname = container.getNodeId.getHost
+         val containerId = container.getId
+ 
+         assert(
+           container.getResource.getMemory >= (workerMemory + YarnAllocationHandler.MEMORY_OVERHEAD))
+ 
+         if (numWorkersRunningNow > maxWorkers) {
+           logInfo("""Ignoring container %s at host %s, since we already have the required number of
+             containers for it.""".format(containerId, workerHostname))
+           releasedContainerList.add(containerId)
+           // reset counter back to old value.
+           numWorkersRunning.decrementAndGet()
+         }
+         else {
+           // Deallocate + allocate can result in reusing id's wrongly - so use a different counter
+           // (workerIdCounter)
+           val workerId = workerIdCounter.incrementAndGet().toString
+           val driverUrl = "akka.tcp://spark@%s:%s/user/%s".format(
+             sparkConf.get("spark.driver.host"), sparkConf.get("spark.driver.port"),
+             CoarseGrainedSchedulerBackend.ACTOR_NAME)
+ 
+           logInfo("launching container on " + containerId + " host " + workerHostname)
+           // Just to be safe, simply remove it from pendingReleaseContainers.
+           // Should not be there, but ..
+           pendingReleaseContainers.remove(containerId)
+ 
+           val rack = YarnAllocationHandler.lookupRack(conf, workerHostname)
+           allocatedHostToContainersMap.synchronized {
+             val containerSet = allocatedHostToContainersMap.getOrElseUpdate(workerHostname,
+               new HashSet[ContainerId]())
+ 
+             containerSet += containerId
+             allocatedContainerToHostMap.put(containerId, workerHostname)
+             if (rack != null) {
+               allocatedRackCount.put(rack, allocatedRackCount.getOrElse(rack, 0) + 1)
+             }
+           }
+ 
+           new Thread(
 -            new WorkerRunnable(container, conf, driverUrl, workerId,
++            new WorkerRunnable(container, conf, sparkConf, driverUrl, workerId,
+               workerHostname, workerMemory, workerCores)
+           ).start()
+         }
+       }
+       logDebug("""
+         Finished processing %d containers.
+         Current number of workers running: %d,
+         releasedContainerList: %s,
+         pendingReleaseContainers: %s
+         """.format(
+           allocatedContainers.size,
+           numWorkersRunning.get(),
+           releasedContainerList,
+           pendingReleaseContainers))
+     }
+ 
+ 
+     val completedContainers = amResp.getCompletedContainersStatuses()
+     if (completedContainers.size > 0){
+       logDebug("Completed %d containers, to-be-released: %s".format(
+         completedContainers.size, releasedContainerList))
+       for (completedContainer <- completedContainers){
+         val containerId = completedContainer.getContainerId
+ 
+         // Was this released by us ? If yes, then simply remove from containerSet and move on.
+         if (pendingReleaseContainers.containsKey(containerId)) {
+           pendingReleaseContainers.remove(containerId)
+         }
+         else {
+           // Simply decrement count - next iteration of ReporterThread will take care of allocating.
+           numWorkersRunning.decrementAndGet()
+           logInfo("Completed container %s (state: %s, exit status: %s)".format(
+             containerId,
+             completedContainer.getState,
+             completedContainer.getExitStatus()))
+           // Hadoop 2.2.X added a ContainerExitStatus we should switch to use
+           // there are some exit status' we shouldn't necessarily count against us, but for
+           // now I think its ok as none of the containers are expected to exit
+           if (completedContainer.getExitStatus() != 0) {
+             logInfo("Container marked as failed: " + containerId)
+             numWorkersFailed.incrementAndGet()
+           }
+         }
+ 
+         allocatedHostToContainersMap.synchronized {
+           if (allocatedContainerToHostMap.containsKey(containerId)) {
+             val host = allocatedContainerToHostMap.get(containerId).getOrElse(null)
+             assert (host != null)
+ 
+             val containerSet = allocatedHostToContainersMap.get(host).getOrElse(null)
+             assert (containerSet != null)
+ 
+             containerSet -= containerId
+             if (containerSet.isEmpty) allocatedHostToContainersMap.remove(host)
+             else allocatedHostToContainersMap.update(host, containerSet)
+ 
+             allocatedContainerToHostMap -= containerId
+ 
+             // Doing this within locked context, sigh ... move to outside ?
+             val rack = YarnAllocationHandler.lookupRack(conf, host)
+             if (rack != null) {
+               val rackCount = allocatedRackCount.getOrElse(rack, 0) - 1
+               if (rackCount > 0) allocatedRackCount.put(rack, rackCount)
+               else allocatedRackCount.remove(rack)
+             }
+           }
+         }
+       }
+       logDebug("""
+         Finished processing %d completed containers.
+         Current number of workers running: %d,
+         releasedContainerList: %s,
+         pendingReleaseContainers: %s
+         """.format(
+           completedContainers.size,
+           numWorkersRunning.get(),
+           releasedContainerList,
+           pendingReleaseContainers))
+     }
+   }
+ 
+   def createRackResourceRequests(hostContainers: List[ResourceRequest]): List[ResourceRequest] = {
+     // First generate modified racks and new set of hosts under it : then issue requests
+     val rackToCounts = new HashMap[String, Int]()
+ 
+     // Within this lock - used to read/write to the rack related maps too.
+     for (container <- hostContainers) {
+       val candidateHost = container.getHostName
+       val candidateNumContainers = container.getNumContainers
+       assert(YarnAllocationHandler.ANY_HOST != candidateHost)
+ 
+       val rack = YarnAllocationHandler.lookupRack(conf, candidateHost)
+       if (rack != null) {
+         var count = rackToCounts.getOrElse(rack, 0)
+         count += candidateNumContainers
+         rackToCounts.put(rack, count)
+       }
+     }
+ 
+     val requestedContainers: ArrayBuffer[ResourceRequest] = 
+       new ArrayBuffer[ResourceRequest](rackToCounts.size)
+     for ((rack, count) <- rackToCounts){
+       requestedContainers += 
+         createResourceRequest(AllocationType.RACK, rack, count, YarnAllocationHandler.PRIORITY)
+     }
+ 
+     requestedContainers.toList
+   }
+ 
+   def allocatedContainersOnHost(host: String): Int = {
+     var retval = 0
+     allocatedHostToContainersMap.synchronized {
+       retval = allocatedHostToContainersMap.getOrElse(host, Set()).size
+     }
+     retval
+   }
+ 
+   def allocatedContainersOnRack(rack: String): Int = {
+     var retval = 0
+     allocatedHostToContainersMap.synchronized {
+       retval = allocatedRackCount.getOrElse(rack, 0)
+     }
+     retval
+   }
+ 
+   private def allocateWorkerResources(numWorkers: Int): AllocateResponse = {
+ 
+     var resourceRequests: List[ResourceRequest] = null
+ 
+       // default.
+     if (numWorkers <= 0 || preferredHostToCount.isEmpty) {
+       logDebug("numWorkers: " + numWorkers + ", host preferences: " + preferredHostToCount.isEmpty)
+       resourceRequests = List(
+         createResourceRequest(AllocationType.ANY, null, numWorkers, YarnAllocationHandler.PRIORITY))
+     }
+     else {
+       // request for all hosts in preferred nodes and for numWorkers - 
+       // candidates.size, request by default allocation policy.
+       val hostContainerRequests: ArrayBuffer[ResourceRequest] = 
+         new ArrayBuffer[ResourceRequest](preferredHostToCount.size)
+       for ((candidateHost, candidateCount) <- preferredHostToCount) {
+         val requiredCount = candidateCount - allocatedContainersOnHost(candidateHost)
+ 
+         if (requiredCount > 0) {
+           hostContainerRequests += createResourceRequest(
+             AllocationType.HOST,
+             candidateHost,
+             requiredCount,
+             YarnAllocationHandler.PRIORITY)
+         }
+       }
+       val rackContainerRequests: List[ResourceRequest] = createRackResourceRequests(
+         hostContainerRequests.toList)
+ 
+       val anyContainerRequests: ResourceRequest = createResourceRequest(
+         AllocationType.ANY,
+         resource = null,
+         numWorkers,
+         YarnAllocationHandler.PRIORITY)
+ 
+       val containerRequests: ArrayBuffer[ResourceRequest] = new ArrayBuffer[ResourceRequest](
+         hostContainerRequests.size + rackContainerRequests.size + 1)
+ 
+       containerRequests ++= hostContainerRequests
+       containerRequests ++= rackContainerRequests
+       containerRequests += anyContainerRequests
+ 
+       resourceRequests = containerRequests.toList
+     }
+ 
+     val req = Records.newRecord(classOf[AllocateRequest])
+     req.setResponseId(lastResponseId.incrementAndGet)
+     req.setApplicationAttemptId(appAttemptId)
+ 
+     req.addAllAsks(resourceRequests)
+ 
+     val releasedContainerList = createReleasedContainerList()
+     req.addAllReleases(releasedContainerList)
+ 
+     if (numWorkers > 0) {
+       logInfo("Allocating %d worker containers with %d of memory each.".format(numWorkers,
+         workerMemory + YarnAllocationHandler.MEMORY_OVERHEAD))
+     }
+     else {
+       logDebug("Empty allocation req ..  release : " + releasedContainerList)
+     }
+ 
+     for (request <- resourceRequests) {
+       logInfo("ResourceRequest (host : %s, num containers: %d, priority = %s , capability : %s)".
+         format(
+           request.getHostName,
+           request.getNumContainers,
+           request.getPriority,
+           request.getCapability))
+     }
+     resourceManager.allocate(req)
+   }
+ 
+ 
+   private def createResourceRequest(
+     requestType: AllocationType.AllocationType, 
+     resource:String,
+     numWorkers: Int,
+     priority: Int): ResourceRequest = {
+ 
+     // If hostname specified, we need atleast two requests - node local and rack local.
+     // There must be a third request - which is ANY : that will be specially handled.
+     requestType match {
+       case AllocationType.HOST => {
+         assert(YarnAllocationHandler.ANY_HOST != resource)
+         val hostname = resource
+         val nodeLocal = createResourceRequestImpl(hostname, numWorkers, priority)
+ 
+         // Add to host->rack mapping
+         YarnAllocationHandler.populateRackInfo(conf, hostname)
+ 
+         nodeLocal
+       }
+       case AllocationType.RACK => {
+         val rack = resource
+         createResourceRequestImpl(rack, numWorkers, priority)
+       }
+       case AllocationType.ANY => createResourceRequestImpl(
+         YarnAllocationHandler.ANY_HOST, numWorkers, priority)
+       case _ => throw new IllegalArgumentException(
+         "Unexpected/unsupported request type: " + requestType)
+     }
+   }
+ 
+   private def createResourceRequestImpl(
+     hostname:String,
+     numWorkers: Int,
+     priority: Int): ResourceRequest = {
+ 
+     val rsrcRequest = Records.newRecord(classOf[ResourceRequest])
+     val memCapability = Records.newRecord(classOf[Resource])
+     // There probably is some overhead here, let's reserve a bit more memory.
+     memCapability.setMemory(workerMemory + YarnAllocationHandler.MEMORY_OVERHEAD)
+     rsrcRequest.setCapability(memCapability)
+ 
+     val pri = Records.newRecord(classOf[Priority])
+     pri.setPriority(priority)
+     rsrcRequest.setPriority(pri)
+ 
+     rsrcRequest.setHostName(hostname)
+ 
+     rsrcRequest.setNumContainers(java.lang.Math.max(numWorkers, 0))
+     rsrcRequest
+   }
+ 
+   def createReleasedContainerList(): ArrayBuffer[ContainerId] = {
+ 
+     val retval = new ArrayBuffer[ContainerId](1)
+     // Iterator on COW list ...
+     for (container <- releasedContainerList.iterator()){
+       retval += container
+     }
+     // Remove from the original list.
+     if (! retval.isEmpty) {
+       releasedContainerList.removeAll(retval)
+       for (v <- retval) pendingReleaseContainers.put(v, true)
+       logInfo("Releasing " + retval.size + " containers. pendingReleaseContainers : " + 
+         pendingReleaseContainers)
+     }
+ 
+     retval
+   }
+ }
+ 
+ object YarnAllocationHandler {
+ 
+   val ANY_HOST = "*"
+   // All requests are issued with same priority : we do not (yet) have any distinction between 
+   // request types (like map/reduce in hadoop for example)
+   val PRIORITY = 1
+ 
+   // Additional memory overhead - in mb
+   val MEMORY_OVERHEAD = 384
+ 
+   // Host to rack map - saved from allocation requests
+   // We are expecting this not to change.
+   // Note that it is possible for this to change : and RM will indicate that to us via update 
+   // response to allocate. But we are punting on handling that for now.
+   private val hostToRack = new ConcurrentHashMap[String, String]()
+   private val rackToHostSet = new ConcurrentHashMap[String, JSet[String]]()
+ 
+ 
+   def newAllocator(
+     conf: Configuration,
+     resourceManager: AMRMProtocol,
+     appAttemptId: ApplicationAttemptId,
+     args: ApplicationMasterArguments,
+     sparkConf: SparkConf): YarnAllocationHandler = {
+ 
+     new YarnAllocationHandler(
+       conf,
+       resourceManager,
+       appAttemptId,
+       args.numWorkers, 
+       args.workerMemory,
+       args.workerCores,
+       Map[String, Int](),
+       Map[String, Int](),
+       sparkConf)
+   }
+ 
+   def newAllocator(
+     conf: Configuration,
+     resourceManager: AMRMProtocol,
+     appAttemptId: ApplicationAttemptId,
+     args: ApplicationMasterArguments,
+     map: collection.Map[String,
+     collection.Set[SplitInfo]],
+     sparkConf: SparkConf): YarnAllocationHandler = {
+ 
+     val (hostToCount, rackToCount) = generateNodeToWeight(conf, map)
+     new YarnAllocationHandler(
+       conf,
+       resourceManager,
+       appAttemptId,
+       args.numWorkers, 
+       args.workerMemory,
+       args.workerCores,
+       hostToCount,
+       rackToCount,
+       sparkConf)
+   }
+ 
+   def newAllocator(
+     conf: Configuration,
+     resourceManager: AMRMProtocol,
+     appAttemptId: ApplicationAttemptId,
+     maxWorkers: Int,
+     workerMemory: Int,
+     workerCores: Int,
+     map: collection.Map[String, collection.Set[SplitInfo]],
+     sparkConf: SparkConf): YarnAllocationHandler = {
+ 
+     val (hostToCount, rackToCount) = generateNodeToWeight(conf, map)
+ 
+     new YarnAllocationHandler(
+       conf,
+       resourceManager,
+       appAttemptId,
+       maxWorkers,
+       workerMemory,
+       workerCores,
+       hostToCount,
+       rackToCount,
+       sparkConf)
+   }
+ 
+   // A simple method to copy the split info map.
+   private def generateNodeToWeight(
+     conf: Configuration,
+     input: collection.Map[String, collection.Set[SplitInfo]]) :
+   // host to count, rack to count
+   (Map[String, Int], Map[String, Int]) = {
+ 
+     if (input == null) return (Map[String, Int](), Map[String, Int]())
+ 
+     val hostToCount = new HashMap[String, Int]
+     val rackToCount = new HashMap[String, Int]
+ 
+     for ((host, splits) <- input) {
+       val hostCount = hostToCount.getOrElse(host, 0)
+       hostToCount.put(host, hostCount + splits.size)
+ 
+       val rack = lookupRack(conf, host)
+       if (rack != null){
+         val rackCount = rackToCount.getOrElse(host, 0)
+         rackToCount.put(host, rackCount + splits.size)
+       }
+     }
+ 
+     (hostToCount.toMap, rackToCount.toMap)
+   }
+ 
+   def lookupRack(conf: Configuration, host: String): String = {
+     if (!hostToRack.contains(host)) populateRackInfo(conf, host)
+     hostToRack.get(host)
+   }
+ 
+   def fetchCachedHostsForRack(rack: String): Option[Set[String]] = {
+     val set = rackToHostSet.get(rack)
+     if (set == null) return None
+ 
+     // No better way to get a Set[String] from JSet ?
+     val convertedSet: collection.mutable.Set[String] = set
+     Some(convertedSet.toSet)
+   }
+ 
+   def populateRackInfo(conf: Configuration, hostname: String) {
+     Utils.checkHost(hostname)
+ 
+     if (!hostToRack.containsKey(hostname)) {
+       // If there are repeated failures to resolve, all to an ignore list ?
+       val rackInfo = RackResolver.resolve(conf, hostname)
+       if (rackInfo != null && rackInfo.getNetworkLocation != null) {
+         val rack = rackInfo.getNetworkLocation
+         hostToRack.put(hostname, rack)
+         if (! rackToHostSet.containsKey(rack)) {
+           rackToHostSet.putIfAbsent(rack,
+             Collections.newSetFromMap(new ConcurrentHashMap[String, JBoolean]()))
+         }
+         rackToHostSet.get(rack).add(hostname)
+ 
+         // TODO(harvey): Figure out this comment...
+         // Since RackResolver caches, we are disabling this for now ...
+       } /* else {
+         // right ? Else we will keep calling rack resolver in case we cant resolve rack info ...
+         hostToRack.put(hostname, null)
+       } */
+     }
+   }
+ }

http://git-wip-us.apache.org/repos/asf/incubator-spark/blob/8ddbd531/yarn/common/src/main/scala/org/apache/spark/deploy/yarn/ClientArguments.scala
----------------------------------------------------------------------
diff --cc yarn/common/src/main/scala/org/apache/spark/deploy/yarn/ClientArguments.scala
index 0000000,7aac232..1419f21
mode 000000,100644..100644
--- a/yarn/common/src/main/scala/org/apache/spark/deploy/yarn/ClientArguments.scala
+++ b/yarn/common/src/main/scala/org/apache/spark/deploy/yarn/ClientArguments.scala
@@@ -1,0 -1,150 +1,150 @@@
+ /*
+  * 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.spark.deploy.yarn
+ 
+ import scala.collection.mutable.{ArrayBuffer, HashMap}
+ 
+ import org.apache.spark.SparkConf
+ import org.apache.spark.scheduler.{InputFormatInfo, SplitInfo}
+ import org.apache.spark.util.IntParam
+ import org.apache.spark.util.MemoryParam
+ 
+ 
+ // TODO: Add code and support for ensuring that yarn resource 'tasks' are location aware !
 -class ClientArguments(val args: Array[String]) {
++class ClientArguments(val args: Array[String], val sparkConf: SparkConf) {
+   var addJars: String = null
+   var files: String = null
+   var archives: String = null
+   var userJar: String = null
+   var userClass: String = null
+   var userArgs: Seq[String] = Seq[String]()
+   var workerMemory = 1024 // MB
+   var workerCores = 1
+   var numWorkers = 2
 -  var amQueue = new SparkConf().get("QUEUE", "default")
++  var amQueue = sparkConf.get("QUEUE", "default")
+   var amMemory: Int = 512 // MB
+   var amClass: String = "org.apache.spark.deploy.yarn.ApplicationMaster"
+   var appName: String = "Spark"
+   // TODO
+   var inputFormatInfo: List[InputFormatInfo] = null
+   // TODO(harvey)
+   var priority = 0
+ 
+   parseArgs(args.toList)
+ 
+   private def parseArgs(inputArgs: List[String]): Unit = {
+     val userArgsBuffer: ArrayBuffer[String] = new ArrayBuffer[String]()
+     val inputFormatMap: HashMap[String, InputFormatInfo] = new HashMap[String, InputFormatInfo]()
+ 
+     var args = inputArgs
+ 
+     while (!args.isEmpty) {
+       args match {
+         case ("--jar") :: value :: tail =>
+           userJar = value
+           args = tail
+ 
+         case ("--class") :: value :: tail =>
+           userClass = value
+           args = tail
+ 
+         case ("--args") :: value :: tail =>
+           userArgsBuffer += value
+           args = tail
+ 
+         case ("--master-class") :: value :: tail =>
+           amClass = value
+           args = tail
+ 
+         case ("--master-memory") :: MemoryParam(value) :: tail =>
+           amMemory = value
+           args = tail
+ 
+         case ("--num-workers") :: IntParam(value) :: tail =>
+           numWorkers = value
+           args = tail
+ 
+         case ("--worker-memory") :: MemoryParam(value) :: tail =>
+           workerMemory = value
+           args = tail
+ 
+         case ("--worker-cores") :: IntParam(value) :: tail =>
+           workerCores = value
+           args = tail
+ 
+         case ("--queue") :: value :: tail =>
+           amQueue = value
+           args = tail
+ 
+         case ("--name") :: value :: tail =>
+           appName = value
+           args = tail
+ 
+         case ("--addJars") :: value :: tail =>
+           addJars = value
+           args = tail
+ 
+         case ("--files") :: value :: tail =>
+           files = value
+           args = tail
+ 
+         case ("--archives") :: value :: tail =>
+           archives = value
+           args = tail
+ 
+         case Nil =>
+           if (userJar == null || userClass == null) {
+             printUsageAndExit(1)
+           }
+ 
+         case _ =>
+           printUsageAndExit(1, args)
+       }
+     }
+ 
+     userArgs = userArgsBuffer.readOnly
+     inputFormatInfo = inputFormatMap.values.toList
+   }
+ 
+ 
+   def printUsageAndExit(exitCode: Int, unknownParam: Any = null) {
+     if (unknownParam != null) {
+       System.err.println("Unknown/unsupported param " + unknownParam)
+     }
+     System.err.println(
+       "Usage: org.apache.spark.deploy.yarn.Client [options] \n" +
+       "Options:\n" +
+       "  --jar JAR_PATH             Path to your application's JAR file (required)\n" +
+       "  --class CLASS_NAME         Name of your application's main class (required)\n" +
+       "  --args ARGS                Arguments to be passed to your application's main class.\n" +
+       "                             Mutliple invocations are possible, each will be passed in order.\n" +
+       "  --num-workers NUM          Number of workers to start (Default: 2)\n" +
+       "  --worker-cores NUM         Number of cores for the workers (Default: 1). This is unsused right now.\n" +
+       "  --master-class CLASS_NAME  Class Name for Master (Default: spark.deploy.yarn.ApplicationMaster)\n" +
+       "  --master-memory MEM        Memory for Master (e.g. 1000M, 2G) (Default: 512 Mb)\n" +
+       "  --worker-memory MEM        Memory per Worker (e.g. 1000M, 2G) (Default: 1G)\n" +
+       "  --name NAME                The name of your application (Default: Spark)\n" +
+       "  --queue QUEUE              The hadoop queue to use for allocation requests (Default: 'default')\n" +
+       "  --addJars jars             Comma separated list of local jars that want SparkContext.addJar to work with.\n" +
+       "  --files files              Comma separated list of files to be distributed with the job.\n" +
+       "  --archives archives        Comma separated list of archives to be distributed with the job."
+       )
+     System.exit(exitCode)
+   }
+ 
+ }

http://git-wip-us.apache.org/repos/asf/incubator-spark/blob/8ddbd531/yarn/common/src/main/scala/org/apache/spark/scheduler/cluster/YarnClientSchedulerBackend.scala
----------------------------------------------------------------------
diff --cc yarn/common/src/main/scala/org/apache/spark/scheduler/cluster/YarnClientSchedulerBackend.scala
index 0000000,4b69f50..324ef46
mode 000000,100644..100644
--- a/yarn/common/src/main/scala/org/apache/spark/scheduler/cluster/YarnClientSchedulerBackend.scala
+++ b/yarn/common/src/main/scala/org/apache/spark/scheduler/cluster/YarnClientSchedulerBackend.scala
@@@ -1,0 -1,110 +1,110 @@@
+ /*
+  * 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.spark.scheduler.cluster
+ 
+ import org.apache.hadoop.yarn.api.records.{ApplicationId, YarnApplicationState}
+ import org.apache.spark.{SparkException, Logging, SparkContext}
+ import org.apache.spark.deploy.yarn.{Client, ClientArguments}
+ import org.apache.spark.scheduler.TaskSchedulerImpl
+ 
+ private[spark] class YarnClientSchedulerBackend(
+     scheduler: TaskSchedulerImpl,
+     sc: SparkContext)
+   extends CoarseGrainedSchedulerBackend(scheduler, sc.env.actorSystem)
+   with Logging {
+ 
+   var client: Client = null
+   var appId: ApplicationId = null
+ 
+   override def start() {
+     super.start()
+ 
+     val defalutWorkerCores = "2"
+     val defalutWorkerMemory = "512m"
+     val defaultWorkerNumber = "1"
+ 
+     val userJar = System.getenv("SPARK_YARN_APP_JAR")
+     var workerCores = System.getenv("SPARK_WORKER_CORES")
+     var workerMemory = System.getenv("SPARK_WORKER_MEMORY")
+     var workerNumber = System.getenv("SPARK_WORKER_INSTANCES")
+ 
+     if (userJar == null)
+       throw new SparkException("env SPARK_YARN_APP_JAR is not set")
+ 
+     if (workerCores == null)
+       workerCores = defalutWorkerCores
+     if (workerMemory == null)
+       workerMemory = defalutWorkerMemory
+     if (workerNumber == null)
+       workerNumber = defaultWorkerNumber
+ 
+     val driverHost = conf.get("spark.driver.host")
+     val driverPort = conf.get("spark.driver.port")
+     val hostport = driverHost + ":" + driverPort
+ 
+     val argsArray = Array[String](
+       "--class", "notused",
+       "--jar", userJar,
+       "--args", hostport,
+       "--worker-memory", workerMemory,
+       "--worker-cores", workerCores,
+       "--num-workers", workerNumber,
+       "--master-class", "org.apache.spark.deploy.yarn.WorkerLauncher"
+     )
+ 
 -    val args = new ClientArguments(argsArray)
 -    client = new Client(args)
++    val args = new ClientArguments(argsArray, conf)
++    client = new Client(args, conf)
+     appId = client.runApp()
+     waitForApp()
+   }
+ 
+   def waitForApp() {
+ 
+     // TODO : need a better way to find out whether the workers are ready or not
+     // maybe by resource usage report?
+     while(true) {
+       val report = client.getApplicationReport(appId)
+ 
+       logInfo("Application report from ASM: \n" +
+         "\t appMasterRpcPort: " + report.getRpcPort() + "\n" +
+         "\t appStartTime: " + report.getStartTime() + "\n" +
+         "\t yarnAppState: " + report.getYarnApplicationState() + "\n"
+       )
+ 
+       // Ready to go, or already gone.
+       val state = report.getYarnApplicationState()
+       if (state == YarnApplicationState.RUNNING) {
+         return
+       } else if (state == YarnApplicationState.FINISHED ||
+         state == YarnApplicationState.FAILED ||
+         state == YarnApplicationState.KILLED) {
+         throw new SparkException("Yarn application already ended," +
+           "might be killed or not able to launch application master.")
+       }
+ 
+       Thread.sleep(1000)
+     }
+   }
+ 
+   override def stop() {
+     super.stop()
+     client.stop()
+     logInfo("Stoped")
+   }
+ 
+ }

http://git-wip-us.apache.org/repos/asf/incubator-spark/blob/8ddbd531/yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
----------------------------------------------------------------------
diff --cc yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
index 0000000,7c32e0a..69ae14c
mode 000000,100644..100644
--- a/yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
+++ b/yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/ApplicationMaster.scala
@@@ -1,0 -1,428 +1,432 @@@
+ /*
+  * 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.spark.deploy.yarn
+ 
+ import java.io.IOException
+ import java.net.Socket
+ import java.util.concurrent.CopyOnWriteArrayList
+ import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}
+ 
+ import scala.collection.JavaConversions._
+ 
+ import org.apache.hadoop.conf.Configuration
+ import org.apache.hadoop.fs.{FileSystem, Path}
+ import org.apache.hadoop.net.NetUtils
+ import org.apache.hadoop.security.UserGroupInformation
+ import org.apache.hadoop.util.ShutdownHookManager
+ import org.apache.hadoop.yarn.api._
+ import org.apache.hadoop.yarn.api.protocolrecords._
+ import org.apache.hadoop.yarn.api.records._
+ import org.apache.hadoop.yarn.client.api.AMRMClient
+ import org.apache.hadoop.yarn.client.api.AMRMClient.ContainerRequest
+ import org.apache.hadoop.yarn.conf.YarnConfiguration
+ import org.apache.hadoop.yarn.ipc.YarnRPC
+ import org.apache.hadoop.yarn.util.{ConverterUtils, Records}
+ 
+ import org.apache.spark.{SparkConf, SparkContext, Logging}
+ import org.apache.spark.util.Utils
+ 
+ 
 -class ApplicationMaster(args: ApplicationMasterArguments, conf: Configuration) extends Logging {
++class ApplicationMaster(args: ApplicationMasterArguments, conf: Configuration,
++                        sparkConf: SparkConf) extends Logging {
+ 
 -  def this(args: ApplicationMasterArguments) = this(args, new Configuration())
++  def this(args: ApplicationMasterArguments, sparkConf: SparkConf) =
++    this(args, new Configuration(), sparkConf)
++
++  def this(args: ApplicationMasterArguments) = this(args, new SparkConf())
+ 
 -  private var rpc: YarnRPC = YarnRPC.create(conf)
+   private val yarnConf: YarnConfiguration = new YarnConfiguration(conf)
+   private var appAttemptId: ApplicationAttemptId = _
+   private var userThread: Thread = _
+   private val fs = FileSystem.get(yarnConf)
+ 
+   private var yarnAllocator: YarnAllocationHandler = _
+   private var isFinished: Boolean = false
+   private var uiAddress: String = _
+   private val maxAppAttempts: Int = conf.getInt(
+     YarnConfiguration.RM_AM_MAX_ATTEMPTS, YarnConfiguration.DEFAULT_RM_AM_MAX_ATTEMPTS)
+   private var isLastAMRetry: Boolean = true
+   private var amClient: AMRMClient[ContainerRequest] = _
+ 
 -  private val sparkConf = new SparkConf()
+   // Default to numWorkers * 2, with minimum of 3
+   private val maxNumWorkerFailures = sparkConf.getInt("spark.yarn.max.worker.failures",
+     math.max(args.numWorkers * 2, 3))
+ 
+   def run() {
+     // Setup the directories so things go to YARN approved directories rather
+     // than user specified and /tmp.
+     System.setProperty("spark.local.dir", getLocalDirs())
+ 
+     // set the web ui port to be ephemeral for yarn so we don't conflict with
+     // other spark processes running on the same box
+     System.setProperty("spark.ui.port", "0")
+ 
+     // Use priority 30 as it's higher then HDFS. It's same priority as MapReduce is using.
+     ShutdownHookManager.get().addShutdownHook(new AppMasterShutdownHook(this), 30)
+ 
+     appAttemptId = getApplicationAttemptId()
+     isLastAMRetry = appAttemptId.getAttemptId() >= maxAppAttempts
+     amClient = AMRMClient.createAMRMClient()
+     amClient.init(yarnConf)
+     amClient.start()
+ 
+     // Workaround until hadoop moves to something which has
+     // https://issues.apache.org/jira/browse/HADOOP-8406 - fixed in (2.0.2-alpha but no 0.23 line)
+     // org.apache.hadoop.io.compress.CompressionCodecFactory.getCodecClasses(conf)
+ 
+     ApplicationMaster.register(this)
+ 
+     // Start the user's JAR
+     userThread = startUserClass()
+ 
+     // This a bit hacky, but we need to wait until the spark.driver.port property has
+     // been set by the Thread executing the user class.
+     waitForSparkContextInitialized()
+ 
+     // Do this after Spark master is up and SparkContext is created so that we can register UI Url.
+     val appMasterResponse: RegisterApplicationMasterResponse = registerApplicationMaster()
+ 
+     // Allocate all containers
+     allocateWorkers()
+ 
+     // Wait for the user class to Finish
+     userThread.join()
+ 
+     System.exit(0)
+   }
+ 
+   /** Get the Yarn approved local directories. */
+   private def getLocalDirs(): String = {
+     // Hadoop 0.23 and 2.x have different Environment variable names for the
+     // local dirs, so lets check both. We assume one of the 2 is set.
+     // LOCAL_DIRS => 2.X, YARN_LOCAL_DIRS => 0.23.X
+     val localDirs = Option(System.getenv("YARN_LOCAL_DIRS"))
+       .getOrElse(Option(System.getenv("LOCAL_DIRS"))
 -        .getOrElse(""))
++      .getOrElse(""))
+ 
+     if (localDirs.isEmpty()) {
+       throw new Exception("Yarn Local dirs can't be empty")
+     }
+     localDirs
+   }
+ 
+   private def getApplicationAttemptId(): ApplicationAttemptId = {
+     val envs = System.getenv()
+     val containerIdString = envs.get(ApplicationConstants.Environment.CONTAINER_ID.name())
+     val containerId = ConverterUtils.toContainerId(containerIdString)
+     val appAttemptId = containerId.getApplicationAttemptId()
+     logInfo("ApplicationAttemptId: " + appAttemptId)
+     appAttemptId
+   }
+ 
+   private def registerApplicationMaster(): RegisterApplicationMasterResponse = {
+     logInfo("Registering the ApplicationMaster")
+     amClient.registerApplicationMaster(Utils.localHostName(), 0, uiAddress)
+   }
+ 
 -  private def startUserClass(): Thread  = {
++  private def startUserClass(): Thread = {
+     logInfo("Starting the user JAR in a separate Thread")
+     val mainMethod = Class.forName(
+       args.userClass,
 -      false /* initialize */,
++      false /* initialize */ ,
+       Thread.currentThread.getContextClassLoader).getMethod("main", classOf[Array[String]])
+     val t = new Thread {
+       override def run() {
+         var successed = false
+         try {
+           // Copy
+           var mainArgs: Array[String] = new Array[String](args.userArgs.size)
+           args.userArgs.copyToArray(mainArgs, 0, args.userArgs.size)
+           mainMethod.invoke(null, mainArgs)
+           // some job script has "System.exit(0)" at the end, for example SparkPi, SparkLR
+           // userThread will stop here unless it has uncaught exception thrown out
+           // It need shutdown hook to set SUCCEEDED
+           successed = true
+         } finally {
+           logDebug("finishing main")
+           isLastAMRetry = true
+           if (successed) {
+             ApplicationMaster.this.finishApplicationMaster(FinalApplicationStatus.SUCCEEDED)
+           } else {
+             ApplicationMaster.this.finishApplicationMaster(FinalApplicationStatus.FAILED)
+           }
+         }
+       }
+     }
+     t.start()
+     t
+   }
+ 
+   // This need to happen before allocateWorkers()
+   private def waitForSparkContextInitialized() {
+     logInfo("Waiting for Spark context initialization")
+     try {
+       var sparkContext: SparkContext = null
+       ApplicationMaster.sparkContextRef.synchronized {
+         var numTries = 0
+         val waitTime = 10000L
+         val maxNumTries = sparkConf.getInt("spark.yarn.applicationMaster.waitTries", 10)
+         while (ApplicationMaster.sparkContextRef.get() == null && numTries < maxNumTries) {
+           logInfo("Waiting for Spark context initialization ... " + numTries)
+           numTries = numTries + 1
+           ApplicationMaster.sparkContextRef.wait(waitTime)
+         }
+         sparkContext = ApplicationMaster.sparkContextRef.get()
+         assert(sparkContext != null || numTries >= maxNumTries)
+ 
+         if (sparkContext != null) {
+           uiAddress = sparkContext.ui.appUIAddress
+           this.yarnAllocator = YarnAllocationHandler.newAllocator(
+             yarnConf,
+             amClient,
+             appAttemptId,
+             args,
+             sparkContext.preferredNodeLocationData,
+             sparkContext.getConf)
+         } else {
+           logWarning("Unable to retrieve SparkContext inspite of waiting for %d, maxNumTries = %d".
+             format(numTries * waitTime, maxNumTries))
+           this.yarnAllocator = YarnAllocationHandler.newAllocator(
+             yarnConf,
+             amClient,
+             appAttemptId,
+             args,
+             sparkContext.getConf)
+         }
+       }
+     } finally {
+       // In case of exceptions, etc - ensure that count is at least ALLOCATOR_LOOP_WAIT_COUNT :
+       // so that the loop (in ApplicationMaster.sparkContextInitialized) breaks.
+       ApplicationMaster.incrementAllocatorLoop(ApplicationMaster.ALLOCATOR_LOOP_WAIT_COUNT)
+     }
+   }
+ 
+   private def allocateWorkers() {
+     try {
+       logInfo("Allocating " + args.numWorkers + " workers.")
+       // Wait until all containers have finished
+       // TODO: This is a bit ugly. Can we make it nicer?
+       // TODO: Handle container failure
+       yarnAllocator.addResourceRequests(args.numWorkers)
+       // Exits the loop if the user thread exits.
+       while (yarnAllocator.getNumWorkersRunning < args.numWorkers && userThread.isAlive) {
+         if (yarnAllocator.getNumWorkersFailed >= maxNumWorkerFailures) {
+           finishApplicationMaster(FinalApplicationStatus.FAILED,
+             "max number of worker failures reached")
+         }
+         yarnAllocator.allocateResources()
+         ApplicationMaster.incrementAllocatorLoop(1)
+         Thread.sleep(100)
+       }
+     } finally {
+       // In case of exceptions, etc - ensure that count is at least ALLOCATOR_LOOP_WAIT_COUNT,
+       // so that the loop in ApplicationMaster#sparkContextInitialized() breaks.
+       ApplicationMaster.incrementAllocatorLoop(ApplicationMaster.ALLOCATOR_LOOP_WAIT_COUNT)
+     }
+     logInfo("All workers have launched.")
+ 
+     // Launch a progress reporter thread, else the app will get killed after expiration
+     // (def: 10mins) timeout.
+     if (userThread.isAlive) {
+       // Ensure that progress is sent before YarnConfiguration.RM_AM_EXPIRY_INTERVAL_MS elapses.
+       val timeoutInterval = yarnConf.getInt(YarnConfiguration.RM_AM_EXPIRY_INTERVAL_MS, 120000)
+ 
+       // we want to be reasonably responsive without causing too many requests to RM.
+       val schedulerInterval =
+         sparkConf.getLong("spark.yarn.scheduler.heartbeat.interval-ms", 5000)
+ 
+ 
+       // must be <= timeoutInterval / 2.
+       val interval = math.min(timeoutInterval / 2, schedulerInterval)
+ 
+       launchReporterThread(interval)
+     }
+   }
+ 
+   private def launchReporterThread(_sleepTime: Long): Thread = {
 -    val sleepTime = if (_sleepTime <= 0 ) 0 else _sleepTime
++    val sleepTime = if (_sleepTime <= 0) 0 else _sleepTime
+ 
+     val t = new Thread {
+       override def run() {
+         while (userThread.isAlive) {
+           if (yarnAllocator.getNumWorkersFailed >= maxNumWorkerFailures) {
+             finishApplicationMaster(FinalApplicationStatus.FAILED,
+               "max number of worker failures reached")
+           }
+           val missingWorkerCount = args.numWorkers - yarnAllocator.getNumWorkersRunning -
+             yarnAllocator.getNumPendingAllocate
+           if (missingWorkerCount > 0) {
+             logInfo("Allocating %d containers to make up for (potentially) lost containers".
+               format(missingWorkerCount))
+             yarnAllocator.addResourceRequests(missingWorkerCount)
+           }
+           sendProgress()
+           Thread.sleep(sleepTime)
+         }
+       }
+     }
+     // Setting to daemon status, though this is usually not a good idea.
+     t.setDaemon(true)
+     t.start()
+     logInfo("Started progress reporter thread - sleep time : " + sleepTime)
+     t
+   }
+ 
+   private def sendProgress() {
+     logDebug("Sending progress")
+     // Simulated with an allocate request with no nodes requested.
+     yarnAllocator.allocateResources()
+   }
+ 
+   /*
+   def printContainers(containers: List[Container]) = {
+     for (container <- containers) {
+       logInfo("Launching shell command on a new container."
+         + ", containerId=" + container.getId()
+         + ", containerNode=" + container.getNodeId().getHost()
+         + ":" + container.getNodeId().getPort()
+         + ", containerNodeURI=" + container.getNodeHttpAddress()
+         + ", containerState" + container.getState()
+         + ", containerResourceMemory"
+         + container.getResource().getMemory())
+     }
+   }
+   */
+ 
+   def finishApplicationMaster(status: FinalApplicationStatus, diagnostics: String = "") {
+     synchronized {
+       if (isFinished) {
+         return
+       }
+       isFinished = true
+     }
+ 
+     logInfo("finishApplicationMaster with " + status)
+     // Set tracking URL to empty since we don't have a history server.
 -    amClient.unregisterApplicationMaster(status, "" /* appMessage */, "" /* appTrackingUrl */)
++    amClient.unregisterApplicationMaster(status, "" /* appMessage */ , "" /* appTrackingUrl */)
+   }
+ 
+   /**
+    * Clean up the staging directory.
+    */
+   private def cleanupStagingDir() {
+     var stagingDirPath: Path = null
+     try {
+       val preserveFiles = sparkConf.get("spark.yarn.preserve.staging.files", "false").toBoolean
+       if (!preserveFiles) {
+         stagingDirPath = new Path(System.getenv("SPARK_YARN_STAGING_DIR"))
+         if (stagingDirPath == null) {
+           logError("Staging directory is null")
+           return
+         }
+         logInfo("Deleting staging directory " + stagingDirPath)
+         fs.delete(stagingDirPath, true)
+       }
+     } catch {
+       case ioe: IOException =>
+         logError("Failed to cleanup staging dir " + stagingDirPath, ioe)
+     }
+   }
+ 
+   // The shutdown hook that runs when a signal is received AND during normal close of the JVM.
+   class AppMasterShutdownHook(appMaster: ApplicationMaster) extends Runnable {
+ 
+     def run() {
+       logInfo("AppMaster received a signal.")
+       // we need to clean up staging dir before HDFS is shut down
+       // make sure we don't delete it until this is the last AM
+       if (appMaster.isLastAMRetry) appMaster.cleanupStagingDir()
+     }
+   }
++
+ }
+ 
+ object ApplicationMaster {
+   // Number of times to wait for the allocator loop to complete.
+   // Each loop iteration waits for 100ms, so maximum of 3 seconds.
+   // This is to ensure that we have reasonable number of containers before we start
+   // TODO: Currently, task to container is computed once (TaskSetManager) - which need not be
+   // optimal as more containers are available. Might need to handle this better.
+   private val ALLOCATOR_LOOP_WAIT_COUNT = 30
+ 
+   private val applicationMasters = new CopyOnWriteArrayList[ApplicationMaster]()
+ 
+   val sparkContextRef: AtomicReference[SparkContext] =
+     new AtomicReference[SparkContext](null /* initialValue */)
+ 
+   val yarnAllocatorLoop: AtomicInteger = new AtomicInteger(0)
+ 
+   def incrementAllocatorLoop(by: Int) {
+     val count = yarnAllocatorLoop.getAndAdd(by)
+     if (count >= ALLOCATOR_LOOP_WAIT_COUNT) {
+       yarnAllocatorLoop.synchronized {
+         // to wake threads off wait ...
+         yarnAllocatorLoop.notifyAll()
+       }
+     }
+   }
+ 
+   def register(master: ApplicationMaster) {
+     applicationMasters.add(master)
+   }
+ 
+   // TODO(harvey): See whether this should be discarded - it isn't used anywhere atm...
+   def sparkContextInitialized(sc: SparkContext): Boolean = {
+     var modified = false
+     sparkContextRef.synchronized {
+       modified = sparkContextRef.compareAndSet(null, sc)
+       sparkContextRef.notifyAll()
+     }
+ 
+     // Add a shutdown hook - as a best case effort in case users do not call sc.stop or do
+     // System.exit.
+     // Should not really have to do this, but it helps YARN to evict resources earlier.
+     // Not to mention, prevent the Client from declaring failure even though we exited properly.
+     // Note that this will unfortunately not properly clean up the staging files because it gets
+     // called too late, after the filesystem is already shutdown.
+     if (modified) {
+       Runtime.getRuntime().addShutdownHook(new Thread with Logging {
+         // This is not only logs, but also ensures that log system is initialized for this instance
+         // when we are actually 'run'-ing.
+         logInfo("Adding shutdown hook for context " + sc)
++
+         override def run() {
+           logInfo("Invoking sc stop from shutdown hook")
+           sc.stop()
+           // Best case ...
+           for (master <- applicationMasters) {
+             master.finishApplicationMaster(FinalApplicationStatus.SUCCEEDED)
+           }
+         }
 -      } )
++      })
+     }
+ 
+     // Wait for initialization to complete and atleast 'some' nodes can get allocated.
+     yarnAllocatorLoop.synchronized {
+       while (yarnAllocatorLoop.get() <= ALLOCATOR_LOOP_WAIT_COUNT) {
+         yarnAllocatorLoop.wait(1000L)
+       }
+     }
+     modified
+   }
+ 
+   def main(argStrings: Array[String]) {
+     val args = new ApplicationMasterArguments(argStrings)
+     new ApplicationMaster(args).run()
+   }
+ }

http://git-wip-us.apache.org/repos/asf/incubator-spark/blob/8ddbd531/yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
----------------------------------------------------------------------
diff --cc yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
index 0000000,a750668..440ad5c
mode 000000,100644..100644
--- a/yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
+++ b/yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
@@@ -1,0 -1,523 +1,525 @@@
+ /*
+  * 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.spark.deploy.yarn
+ 
+ import java.net.{InetAddress, UnknownHostException, URI}
+ import java.nio.ByteBuffer
+ 
+ import scala.collection.JavaConversions._
+ import scala.collection.mutable.HashMap
+ import scala.collection.mutable.Map
+ 
+ import org.apache.hadoop.conf.Configuration
+ import org.apache.hadoop.fs.{FileContext, FileStatus, FileSystem, Path, FileUtil}
+ import org.apache.hadoop.fs.permission.FsPermission;
+ import org.apache.hadoop.io.DataOutputBuffer
+ import org.apache.hadoop.mapred.Master
+ import org.apache.hadoop.net.NetUtils
+ import org.apache.hadoop.security.UserGroupInformation
+ import org.apache.hadoop.yarn.api._
+ import org.apache.hadoop.yarn.api.ApplicationConstants.Environment
+ import org.apache.hadoop.yarn.api.protocolrecords._
+ import org.apache.hadoop.yarn.api.records._
+ import org.apache.hadoop.yarn.client.api.impl.YarnClientImpl
+ import org.apache.hadoop.yarn.conf.YarnConfiguration
+ import org.apache.hadoop.yarn.ipc.YarnRPC
+ import org.apache.hadoop.yarn.util.{Apps, Records}
+ 
+ import org.apache.spark.{Logging, SparkConf}
+ import org.apache.spark.util.Utils
+ import org.apache.spark.deploy.SparkHadoopUtil
+ 
+ 
+ /**
+  * The entry point (starting in Client#main() and Client#run()) for launching Spark on YARN. The
+  * Client submits an application to the global ResourceManager to launch Spark's ApplicationMaster,
+  * which will launch a Spark master process and negotiate resources throughout its duration.
+  */
 -class Client(conf: Configuration, args: ClientArguments) extends YarnClientImpl with Logging {
++class Client(args: ClientArguments, conf: Configuration, sparkConf: SparkConf)
++  extends YarnClientImpl with Logging {
++
++  def this(args: ClientArguments, sparkConf: SparkConf) =
++    this(args, new Configuration(), sparkConf)
++
++  def this(args: ClientArguments) = this(args, new SparkConf())
+ 
+   var rpc: YarnRPC = YarnRPC.create(conf)
+   val yarnConf: YarnConfiguration = new YarnConfiguration(conf)
+   val credentials = UserGroupInformation.getCurrentUser().getCredentials()
+   private val SPARK_STAGING: String = ".sparkStaging"
+   private val distCacheMgr = new ClientDistributedCacheManager()
 -  private val sparkConf = new SparkConf
 -
+ 
+   // Staging directory is private! -> rwx--------
+   val STAGING_DIR_PERMISSION: FsPermission = FsPermission.createImmutable(0700: Short)
+   // App files are world-wide readable and owner writable -> rw-r--r--
+   val APP_FILE_PERMISSION: FsPermission = FsPermission.createImmutable(0644: Short)
+ 
 -  def this(args: ClientArguments) = this(new Configuration(), args)
 -
+   def runApp(): ApplicationId = {
+     validateArgs()
+     // Initialize and start the client service.
+     init(yarnConf)
+     start()
+ 
+     // Log details about this YARN cluster (e.g, the number of slave machines/NodeManagers).
+     logClusterResourceDetails()
+ 
+     // Prepare to submit a request to the ResourcManager (specifically its ApplicationsManager (ASM)
+     // interface).
+ 
+     // Get a new client application.
+     val newApp = super.createApplication()
+     val newAppResponse = newApp.getNewApplicationResponse()
+     val appId = newAppResponse.getApplicationId()
+ 
+     verifyClusterResources(newAppResponse)
+ 
+     // Set up resource and environment variables.
+     val appStagingDir = getAppStagingDir(appId)
+     val localResources = prepareLocalResources(appStagingDir)
+     val launchEnv = setupLaunchEnv(localResources, appStagingDir)
+     val amContainer = createContainerLaunchContext(newAppResponse, localResources, launchEnv)
+ 
+     // Set up an application submission context.
+     val appContext = newApp.getApplicationSubmissionContext()
+     appContext.setApplicationName(args.appName)
+     appContext.setQueue(args.amQueue)
+     appContext.setAMContainerSpec(amContainer)
+ 
+     // Memory for the ApplicationMaster.
+     val memoryResource = Records.newRecord(classOf[Resource]).asInstanceOf[Resource]
+     memoryResource.setMemory(args.amMemory + YarnAllocationHandler.MEMORY_OVERHEAD)
+     appContext.setResource(memoryResource)
+ 
+     // Finally, submit and monitor the application.
+     submitApp(appContext)
+     appId
+   }
+ 
+   def run() {
+     val appId = runApp()
+     monitorApplication(appId)
+     System.exit(0)
+   }
+ 
+   // TODO(harvey): This could just go in ClientArguments.
+   def validateArgs() = {
+     Map(
+       (System.getenv("SPARK_JAR") == null) -> "Error: You must set SPARK_JAR environment variable!",
+       (args.userJar == null) -> "Error: You must specify a user jar!",
+       (args.userClass == null) -> "Error: You must specify a user class!",
+       (args.numWorkers <= 0) -> "Error: You must specify atleast 1 worker!",
+       (args.amMemory <= YarnAllocationHandler.MEMORY_OVERHEAD) -> ("Error: AM memory size must be" +
+         "greater than: " + YarnAllocationHandler.MEMORY_OVERHEAD),
+       (args.workerMemory <= YarnAllocationHandler.MEMORY_OVERHEAD) -> ("Error: Worker memory size" +
+         "must be greater than: " + YarnAllocationHandler.MEMORY_OVERHEAD.toString)
+     ).foreach { case(cond, errStr) =>
+       if (cond) {
+         logError(errStr)
+         args.printUsageAndExit(1)
+       }
+     }
+   }
+ 
+   def getAppStagingDir(appId: ApplicationId): String = {
+     SPARK_STAGING + Path.SEPARATOR + appId.toString() + Path.SEPARATOR
+   }
+ 
+   def logClusterResourceDetails() {
+     val clusterMetrics: YarnClusterMetrics = super.getYarnClusterMetrics
+     logInfo("Got Cluster metric info from ApplicationsManager (ASM), number of NodeManagers: " +
+       clusterMetrics.getNumNodeManagers)
+ 
+     val queueInfo: QueueInfo = super.getQueueInfo(args.amQueue)
 -    logInfo("""Queue info ... queueName: %s, queueCurrentCapacity: %s, queueMaxCapacity: %s,
++    logInfo( """Queue info ... queueName: %s, queueCurrentCapacity: %s, queueMaxCapacity: %s,
+       queueApplicationCount = %s, queueChildQueueCount = %s""".format(
+         queueInfo.getQueueName,
+         queueInfo.getCurrentCapacity,
+         queueInfo.getMaximumCapacity,
+         queueInfo.getApplications.size,
+         queueInfo.getChildQueues.size))
+   }
+ 
+   def verifyClusterResources(app: GetNewApplicationResponse) = {
+     val maxMem = app.getMaximumResourceCapability().getMemory()
+     logInfo("Max mem capabililty of a single resource in this cluster " + maxMem)
+ 
+     // If we have requested more then the clusters max for a single resource then exit.
+     if (args.workerMemory > maxMem) {
+       logError("Required worker memory (%d MB), is above the max threshold (%d MB) of this cluster.".
+         format(args.workerMemory, maxMem))
+       System.exit(1)
+     }
+     val amMem = args.amMemory + YarnAllocationHandler.MEMORY_OVERHEAD
+     if (amMem > maxMem) {
+       logError("Required AM memory (%d) is above the max threshold (%d) of this cluster".
+         format(args.amMemory, maxMem))
+       System.exit(1)
+     }
+ 
+     // We could add checks to make sure the entire cluster has enough resources but that involves
+     // getting all the node reports and computing ourselves.
+   }
+ 
+   /** See if two file systems are the same or not. */
+   private def compareFs(srcFs: FileSystem, destFs: FileSystem): Boolean = {
+     val srcUri = srcFs.getUri()
+     val dstUri = destFs.getUri()
+     if (srcUri.getScheme() == null) {
+       return false
+     }
+     if (!srcUri.getScheme().equals(dstUri.getScheme())) {
+       return false
+     }
+     var srcHost = srcUri.getHost()
+     var dstHost = dstUri.getHost()
+     if ((srcHost != null) && (dstHost != null)) {
+       try {
+         srcHost = InetAddress.getByName(srcHost).getCanonicalHostName()
+         dstHost = InetAddress.getByName(dstHost).getCanonicalHostName()
+       } catch {
+         case e: UnknownHostException =>
+           return false
+       }
+       if (!srcHost.equals(dstHost)) {
+         return false
+       }
+     } else if (srcHost == null && dstHost != null) {
+       return false
+     } else if (srcHost != null && dstHost == null) {
+       return false
+     }
+     //check for ports
+     if (srcUri.getPort() != dstUri.getPort()) {
+       return false
+     }
+     return true
+   }
+ 
+   /** Copy the file into HDFS if needed. */
+   private def copyRemoteFile(
+       dstDir: Path,
+       originalPath: Path,
+       replication: Short,
+       setPerms: Boolean = false): Path = {
+     val fs = FileSystem.get(conf)
+     val remoteFs = originalPath.getFileSystem(conf)
+     var newPath = originalPath
+     if (! compareFs(remoteFs, fs)) {
+       newPath = new Path(dstDir, originalPath.getName())
+       logInfo("Uploading " + originalPath + " to " + newPath)
+       FileUtil.copy(remoteFs, originalPath, fs, newPath, false, conf)
+       fs.setReplication(newPath, replication)
+       if (setPerms) fs.setPermission(newPath, new FsPermission(APP_FILE_PERMISSION))
+     }
+     // Resolve any symlinks in the URI path so using a "current" symlink to point to a specific
+     // version shows the specific version in the distributed cache configuration
+     val qualPath = fs.makeQualified(newPath)
+     val fc = FileContext.getFileContext(qualPath.toUri(), conf)
+     val destPath = fc.resolvePath(qualPath)
+     destPath
+   }
+ 
+   def prepareLocalResources(appStagingDir: String): HashMap[String, LocalResource] = {
+     logInfo("Preparing Local resources")
+     // Upload Spark and the application JAR to the remote file system if necessary. Add them as
+     // local resources to the application master.
+     val fs = FileSystem.get(conf)
+ 
+     val delegTokenRenewer = Master.getMasterPrincipal(conf)
+     if (UserGroupInformation.isSecurityEnabled()) {
+       if (delegTokenRenewer == null || delegTokenRenewer.length() == 0) {
+         logError("Can't get Master Kerberos principal for use as renewer")
+         System.exit(1)
+       }
+     }
+     val dst = new Path(fs.getHomeDirectory(), appStagingDir)
+     val replication = sparkConf.getInt("spark.yarn.submit.file.replication", 3).toShort
+ 
+     if (UserGroupInformation.isSecurityEnabled()) {
+       val dstFs = dst.getFileSystem(conf)
+       dstFs.addDelegationTokens(delegTokenRenewer, credentials)
+     }
+ 
+     val localResources = HashMap[String, LocalResource]()
+     FileSystem.mkdirs(fs, dst, new FsPermission(STAGING_DIR_PERMISSION))
+ 
+     val statCache: Map[URI, FileStatus] = HashMap[URI, FileStatus]()
+ 
+     Map(
+       Client.SPARK_JAR -> System.getenv("SPARK_JAR"), Client.APP_JAR -> args.userJar,
+       Client.LOG4J_PROP -> System.getenv("SPARK_LOG4J_CONF")
+     ).foreach { case(destName, _localPath) =>
+       val localPath: String = if (_localPath != null) _localPath.trim() else ""
+       if (! localPath.isEmpty()) {
+         var localURI = new URI(localPath)
+         // If not specified assume these are in the local filesystem to keep behavior like Hadoop
+         if (localURI.getScheme() == null) {
+           localURI = new URI(FileSystem.getLocal(conf).makeQualified(new Path(localPath)).toString)
+         }
+         val setPermissions = if (destName.equals(Client.APP_JAR)) true else false
+         val destPath = copyRemoteFile(dst, new Path(localURI), replication, setPermissions)
+         distCacheMgr.addResource(fs, conf, destPath, localResources, LocalResourceType.FILE,
+           destName, statCache)
+       }
+     }
+ 
+     // Handle jars local to the ApplicationMaster.
+     if ((args.addJars != null) && (!args.addJars.isEmpty())){
+       args.addJars.split(',').foreach { case file: String =>
+         val localURI = new URI(file.trim())
+         val localPath = new Path(localURI)
+         val linkname = Option(localURI.getFragment()).getOrElse(localPath.getName())
+         val destPath = copyRemoteFile(dst, localPath, replication)
+         // Only add the resource to the Spark ApplicationMaster.
+         val appMasterOnly = true
+         distCacheMgr.addResource(fs, conf, destPath, localResources, LocalResourceType.FILE,
+           linkname, statCache, appMasterOnly)
+       }
+     }
+ 
+     // Handle any distributed cache files
+     if ((args.files != null) && (!args.files.isEmpty())){
+       args.files.split(',').foreach { case file: String =>
+         val localURI = new URI(file.trim())
+         val localPath = new Path(localURI)
+         val linkname = Option(localURI.getFragment()).getOrElse(localPath.getName())
+         val destPath = copyRemoteFile(dst, localPath, replication)
+         distCacheMgr.addResource(fs, conf, destPath, localResources, LocalResourceType.FILE,
+           linkname, statCache)
+       }
+     }
+ 
+     // Handle any distributed cache archives
+     if ((args.archives != null) && (!args.archives.isEmpty())) {
+       args.archives.split(',').foreach { case file:String =>
+         val localURI = new URI(file.trim())
+         val localPath = new Path(localURI)
+         val linkname = Option(localURI.getFragment()).getOrElse(localPath.getName())
+         val destPath = copyRemoteFile(dst, localPath, replication)
+         distCacheMgr.addResource(fs, conf, destPath, localResources, LocalResourceType.ARCHIVE,
+           linkname, statCache)
+       }
+     }
+ 
+     UserGroupInformation.getCurrentUser().addCredentials(credentials)
+     localResources
+   }
+ 
+   def setupLaunchEnv(
+       localResources: HashMap[String, LocalResource],
+       stagingDir: String): HashMap[String, String] = {
+     logInfo("Setting up the launch environment")
+     val log4jConfLocalRes = localResources.getOrElse(Client.LOG4J_PROP, null)
+ 
+     val env = new HashMap[String, String]()
+ 
 -    Client.populateClasspath(yarnConf, log4jConfLocalRes != null, env)
++    Client.populateClasspath(yarnConf, sparkConf, log4jConfLocalRes != null, env)
+     env("SPARK_YARN_MODE") = "true"
+     env("SPARK_YARN_STAGING_DIR") = stagingDir
+ 
+     // Set the environment variables to be passed on to the Workers.
+     distCacheMgr.setDistFilesEnv(env)
+     distCacheMgr.setDistArchivesEnv(env)
+ 
+     // Allow users to specify some environment variables.
+     Apps.setEnvFromInputString(env, System.getenv("SPARK_YARN_USER_ENV"))
+ 
+     // Add each SPARK_* key to the environment.
+     System.getenv().filterKeys(_.startsWith("SPARK")).foreach { case (k,v) => env(k) = v }
+ 
+     env
+   }
+ 
+   def userArgsToString(clientArgs: ClientArguments): String = {
+     val prefix = " --args "
+     val args = clientArgs.userArgs
+     val retval = new StringBuilder()
 -    for (arg <- args){
++    for (arg <- args) {
+       retval.append(prefix).append(" '").append(arg).append("' ")
+     }
+     retval.toString
+   }
+ 
+   def createContainerLaunchContext(
+       newApp: GetNewApplicationResponse,
+       localResources: HashMap[String, LocalResource],
+       env: HashMap[String, String]): ContainerLaunchContext = {
+     logInfo("Setting up container launch context")
+     val amContainer = Records.newRecord(classOf[ContainerLaunchContext])
+     amContainer.setLocalResources(localResources)
+     amContainer.setEnvironment(env)
+ 
+     // TODO: Need a replacement for the following code to fix -Xmx?
+     // val minResMemory: Int = newApp.getMinimumResourceCapability().getMemory()
+     // var amMemory = ((args.amMemory / minResMemory) * minResMemory) +
+     //  ((if ((args.amMemory % minResMemory) == 0) 0 else minResMemory) -
+     //    YarnAllocationHandler.MEMORY_OVERHEAD)
+ 
+     // Extra options for the JVM
+     var JAVA_OPTS = ""
+ 
+     // Add Xmx for AM memory
+     JAVA_OPTS += "-Xmx" + args.amMemory + "m"
+ 
+     val tmpDir = new Path(Environment.PWD.$(), YarnConfiguration.DEFAULT_CONTAINER_TEMP_DIR)
+     JAVA_OPTS += " -Djava.io.tmpdir=" + tmpDir
+ 
+     // TODO: Remove once cpuset version is pushed out.
+     // The context is, default gc for server class machines ends up using all cores to do gc -
+     // hence if there are multiple containers in same node, Spark GC affects all other containers'
+     // performance (which can be that of other Spark containers)
+     // Instead of using this, rely on cpusets by YARN to enforce "proper" Spark behavior in
+     // multi-tenant environments. Not sure how default Java GC behaves if it is limited to subset
+     // of cores on a node.
+     val useConcurrentAndIncrementalGC = env.isDefinedAt("SPARK_USE_CONC_INCR_GC") &&
+       java.lang.Boolean.parseBoolean(env("SPARK_USE_CONC_INCR_GC"))
+     if (useConcurrentAndIncrementalGC) {
+       // In our expts, using (default) throughput collector has severe perf ramifications in
+       // multi-tenant machines
+       JAVA_OPTS += " -XX:+UseConcMarkSweepGC "
+       JAVA_OPTS += " -XX:+CMSIncrementalMode "
+       JAVA_OPTS += " -XX:+CMSIncrementalPacing "
+       JAVA_OPTS += " -XX:CMSIncrementalDutyCycleMin=0 "
+       JAVA_OPTS += " -XX:CMSIncrementalDutyCycle=10 "
+     }
+ 
+     if (env.isDefinedAt("SPARK_JAVA_OPTS")) {
+       JAVA_OPTS += " " + env("SPARK_JAVA_OPTS")
+     }
+ 
+     // Command for the ApplicationMaster
+     var javaCommand = "java"
+     val javaHome = System.getenv("JAVA_HOME")
+     if ((javaHome != null && !javaHome.isEmpty()) || env.isDefinedAt("JAVA_HOME")) {
+       javaCommand = Environment.JAVA_HOME.$() + "/bin/java"
+     }
+ 
+     val commands = List[String](
+       javaCommand +
+       " -server " +
+       JAVA_OPTS +
+       " " + args.amClass +
+       " --class " + args.userClass +
+       " --jar " + args.userJar +
+       userArgsToString(args) +
+       " --worker-memory " + args.workerMemory +
+       " --worker-cores " + args.workerCores +
+       " --num-workers " + args.numWorkers +
+       " 1> " + ApplicationConstants.LOG_DIR_EXPANSION_VAR + "/stdout" +
+       " 2> " + ApplicationConstants.LOG_DIR_EXPANSION_VAR + "/stderr")
+ 
+     logInfo("Command for starting the Spark ApplicationMaster: " + commands(0))
+     amContainer.setCommands(commands)
+ 
+     // Setup security tokens.
+     val dob = new DataOutputBuffer()
+     credentials.writeTokenStorageToStream(dob)
+     amContainer.setTokens(ByteBuffer.wrap(dob.getData()))
+ 
+     amContainer
+   }
+ 
+   def submitApp(appContext: ApplicationSubmissionContext) = {
+     // Submit the application to the applications manager.
+     logInfo("Submitting application to ASM")
+     super.submitApplication(appContext)
+   }
+ 
+   def monitorApplication(appId: ApplicationId): Boolean = {
+     val interval = sparkConf.getLong("spark.yarn.report.interval", 1000)
+ 
+     while (true) {
+       Thread.sleep(interval)
+       val report = super.getApplicationReport(appId)
+ 
+       logInfo("Application report from ASM: \n" +
+         "\t application identifier: " + appId.toString() + "\n" +
+         "\t appId: " + appId.getId() + "\n" +
+         "\t clientToAMToken: " + report.getClientToAMToken() + "\n" +
+         "\t appDiagnostics: " + report.getDiagnostics() + "\n" +
+         "\t appMasterHost: " + report.getHost() + "\n" +
+         "\t appQueue: " + report.getQueue() + "\n" +
+         "\t appMasterRpcPort: " + report.getRpcPort() + "\n" +
+         "\t appStartTime: " + report.getStartTime() + "\n" +
+         "\t yarnAppState: " + report.getYarnApplicationState() + "\n" +
+         "\t distributedFinalState: " + report.getFinalApplicationStatus() + "\n" +
+         "\t appTrackingUrl: " + report.getTrackingUrl() + "\n" +
+         "\t appUser: " + report.getUser()
+       )
+ 
+       val state = report.getYarnApplicationState()
+       val dsStatus = report.getFinalApplicationStatus()
+       if (state == YarnApplicationState.FINISHED ||
+         state == YarnApplicationState.FAILED ||
+         state == YarnApplicationState.KILLED) {
+         return true
+       }
+     }
+     true
+   }
+ }
+ 
+ object Client {
+   val SPARK_JAR: String = "spark.jar"
+   val APP_JAR: String = "app.jar"
+   val LOG4J_PROP: String = "log4j.properties"
+ 
+   def main(argStrings: Array[String]) {
+     // Set an env variable indicating we are running in YARN mode.
+     // Note: anything env variable with SPARK_ prefix gets propagated to all (remote) processes -
+     // see Client#setupLaunchEnv().
+     System.setProperty("SPARK_YARN_MODE", "true")
++    val sparkConf = new SparkConf()
++    val args = new ClientArguments(argStrings, sparkConf)
+ 
 -    val args = new ClientArguments(argStrings)
 -
 -    (new Client(args)).run()
++    new Client(args, sparkConf).run()
+   }
+ 
+   // Based on code from org.apache.hadoop.mapreduce.v2.util.MRApps
+   def populateHadoopClasspath(conf: Configuration, env: HashMap[String, String]) {
+     for (c <- conf.getStrings(YarnConfiguration.YARN_APPLICATION_CLASSPATH)) {
+       Apps.addToEnvironment(env, Environment.CLASSPATH.name, c.trim)
+     }
+   }
+ 
 -  def populateClasspath(conf: Configuration, addLog4j: Boolean, env: HashMap[String, String]) {
++  def populateClasspath(conf: Configuration, sparkConf: SparkConf, addLog4j: Boolean, env: HashMap[String, String]) {
+     Apps.addToEnvironment(env, Environment.CLASSPATH.name, Environment.PWD.$())
+     // If log4j present, ensure ours overrides all others
+     if (addLog4j) {
+       Apps.addToEnvironment(env, Environment.CLASSPATH.name, Environment.PWD.$() +
+         Path.SEPARATOR + LOG4J_PROP)
+     }
+     // Normally the users app.jar is last in case conflicts with spark jars
 -    val userClasspathFirst = new SparkConf().get("spark.yarn.user.classpath.first", "false")
++    val userClasspathFirst = sparkConf.get("spark.yarn.user.classpath.first", "false")
+       .toBoolean
+     if (userClasspathFirst) {
+       Apps.addToEnvironment(env, Environment.CLASSPATH.name, Environment.PWD.$() +
+         Path.SEPARATOR + APP_JAR)
+     }
+     Apps.addToEnvironment(env, Environment.CLASSPATH.name, Environment.PWD.$() +
+       Path.SEPARATOR + SPARK_JAR)
+     Client.populateHadoopClasspath(conf, env)
+ 
+     if (!userClasspathFirst) {
+       Apps.addToEnvironment(env, Environment.CLASSPATH.name, Environment.PWD.$() +
+         Path.SEPARATOR + APP_JAR)
+     }
+     Apps.addToEnvironment(env, Environment.CLASSPATH.name, Environment.PWD.$() +
+       Path.SEPARATOR + "*")
+   }
+ }