You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@samza.apache.org by ni...@apache.org on 2015/09/26 01:51:07 UTC
[2/3] samza git commit: SAMZA-619: SamzaAppMaster changes to enable
host-affinity
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/main/resources/scalate/WEB-INF/views/index.scaml
----------------------------------------------------------------------
diff --git a/samza-yarn/src/main/resources/scalate/WEB-INF/views/index.scaml b/samza-yarn/src/main/resources/scalate/WEB-INF/views/index.scaml
index ce2145a..2d16fe0 100644
--- a/samza-yarn/src/main/resources/scalate/WEB-INF/views/index.scaml
+++ b/samza-yarn/src/main/resources/scalate/WEB-INF/views/index.scaml
@@ -15,7 +15,7 @@
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
--@ val state: org.apache.samza.job.yarn.SamzaAppMasterState
+-@ val state: org.apache.samza.job.yarn.SamzaAppState
-@ val config: scala.collection.immutable.TreeMap[String, String]
-@ val rmHttpAddress: String
-@ val jobName: String = config.get("job.name").getOrElse("MISSING JOB NAME")
@@ -88,7 +88,7 @@
%tr
%td.key Application master container
%td
- %a(target="_blank" href="http://#{state.nodeHost}:#{state.nodeHttpPort.toString}/node/containerlogs/#{state.containerId.toString}/#{username}")= state.containerId.toString
+ %a(target="_blank" href="http://#{state.nodeHost}:#{state.nodeHttpPort.toString}/node/containerlogs/#{state.amContainerId.toString}/#{username}")= state.amContainerId.toString
%tr
%td.key JMX server url
%td= state.jmxUrl
@@ -148,9 +148,6 @@
%td.key Total
%td= state.containerCount.toString
%tr
- %td.key Unclaimed
- %td= state.unclaimedContainers.size.toString
- %tr
%td.key Finished
%td= state.finishedContainers.size.toString
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/main/scala/org/apache/samza/config/YarnConfig.scala
----------------------------------------------------------------------
diff --git a/samza-yarn/src/main/scala/org/apache/samza/config/YarnConfig.scala b/samza-yarn/src/main/scala/org/apache/samza/config/YarnConfig.scala
deleted file mode 100644
index 1f51551..0000000
--- a/samza-yarn/src/main/scala/org/apache/samza/config/YarnConfig.scala
+++ /dev/null
@@ -1,58 +0,0 @@
-/*
- * 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.samza.config
-
-object YarnConfig {
- // yarn job config
- val PACKAGE_PATH = "yarn.package.path"
- val CONTAINER_MAX_MEMORY_MB = "yarn.container.memory.mb"
- val CONTAINER_MAX_CPU_CORES = "yarn.container.cpu.cores"
- val CONTAINER_RETRY_COUNT = "yarn.container.retry.count"
- val CONTAINER_RETRY_WINDOW_MS = "yarn.container.retry.window.ms"
- val AM_JVM_OPTIONS = "yarn.am.opts"
- val AM_JMX_ENABLED = "yarn.am.jmx.enabled"
- val AM_CONTAINER_MAX_MEMORY_MB = "yarn.am.container.memory.mb"
- val AM_POLL_INTERVAL_MS = "yarn.am.poll.interval.ms"
- val AM_JAVA_HOME = "yarn.am.java.home"
-
- implicit def Config2Yarn(config: Config) = new YarnConfig(config)
-}
-
-class YarnConfig(config: Config) extends ScalaMapConfig(config) {
- def getContainerMaxMemoryMb: Option[Int] = getOption(YarnConfig.CONTAINER_MAX_MEMORY_MB).map(_.toInt)
-
- def getContainerMaxCpuCores: Option[Int] = getOption(YarnConfig.CONTAINER_MAX_CPU_CORES).map(_.toInt)
-
- def getContainerRetryCount: Option[Int] = getOption(YarnConfig.CONTAINER_RETRY_COUNT).map(_.toInt)
-
- def getContainerRetryWindowMs: Option[Int] = getOption(YarnConfig.CONTAINER_RETRY_WINDOW_MS).map(_.toInt)
-
- def getPackagePath = getOption(YarnConfig.PACKAGE_PATH)
-
- def getAmOpts = getOption(YarnConfig.AM_JVM_OPTIONS)
-
- def getAMContainerMaxMemoryMb: Option[Int] = getOption(YarnConfig.AM_CONTAINER_MAX_MEMORY_MB).map(_.toInt)
-
- def getAMPollIntervalMs: Option[Int] = getOption(YarnConfig.AM_POLL_INTERVAL_MS).map(_.toInt)
-
- def getJmxServerEnabled = getBoolean(YarnConfig.AM_JMX_ENABLED, true)
-
- def getAMJavaHome = getOption(YarnConfig.AM_JAVA_HOME)
-}
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMaster.scala
----------------------------------------------------------------------
diff --git a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMaster.scala b/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMaster.scala
index af42c6a..80deb3b 100644
--- a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMaster.scala
+++ b/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMaster.scala
@@ -31,9 +31,6 @@ import org.apache.samza.config.Config
import org.apache.samza.config.ShellCommandConfig
import org.apache.samza.config.JobConfig.Config2Job
import org.apache.samza.config.YarnConfig
-import org.apache.samza.config.YarnConfig.Config2Yarn
-import org.apache.samza.job.yarn.SamzaAppMasterTaskManager.DEFAULT_CONTAINER_MEM
-import org.apache.samza.job.yarn.SamzaAppMasterTaskManager.DEFAULT_CPU_CORES
import org.apache.samza.metrics.JmxServer
import org.apache.samza.metrics.MetricsRegistryMap
import org.apache.samza.util.hadoop.HttpFileSystem
@@ -76,21 +73,23 @@ object SamzaAppMaster extends Logging with AMRMClientAsync.CallbackHandler {
val registry = new MetricsRegistryMap
val jobCoordinator = JobCoordinator(coordinatorSystemConfig, registry)
val config = jobCoordinator.jobModel.getConfig
+ val yarnConfig = new YarnConfig(config)
info("got config: %s" format coordinatorSystemConfig)
putMDC("jobName", config.getName.getOrElse(throw new SamzaException("can not find the job name")))
putMDC("jobId", config.getJobId.getOrElse("1"))
val hConfig = new YarnConfiguration
hConfig.set("fs.http.impl", classOf[HttpFileSystem].getName)
- val interval = config.getAMPollIntervalMs.getOrElse(DEFAULT_POLL_INTERVAL_MS)
+ val interval = yarnConfig.getAMPollIntervalMs
val amClient = AMRMClientAsync.createAMRMClientAsync[ContainerRequest](interval, this)
val clientHelper = new ClientHelper(hConfig)
- val containerMem = config.getContainerMaxMemoryMb.getOrElse(DEFAULT_CONTAINER_MEM)
- val containerCpu = config.getContainerMaxCpuCores.getOrElse(DEFAULT_CPU_CORES)
- val jmxServer = if (new YarnConfig(config).getJmxServerEnabled) Some(new JmxServer()) else None
+ val containerMem = yarnConfig.getContainerMaxMemoryMb
+ val containerCpu = yarnConfig.getContainerMaxCpuCores
+ val jmxServer = if (yarnConfig.getJmxServerEnabled) Some(new JmxServer()) else None
try {
// wire up all of the yarn event listeners
- val state = new SamzaAppMasterState(jobCoordinator, -1, containerId, nodeHostString, nodePortString.toInt, nodeHttpPortString.toInt)
+ val state = new SamzaAppState(jobCoordinator, -1, containerId, nodeHostString, nodePortString.toInt, nodeHttpPortString.toInt)
+
if (jmxServer.isDefined) {
state.jmxUrl = jmxServer.get.getJmxUrl
state.jmxTunnelingUrl = jmxServer.get.getTunnelingJmxUrl
@@ -99,8 +98,9 @@ object SamzaAppMaster extends Logging with AMRMClientAsync.CallbackHandler {
val service = new SamzaAppMasterService(config, state, registry, clientHelper)
val lifecycle = new SamzaAppMasterLifecycle(containerMem, containerCpu, state, amClient)
val metrics = new SamzaAppMasterMetrics(config, state, registry)
- val am = new SamzaAppMasterTaskManager({ System.currentTimeMillis }, config, state, amClient, hConfig)
- listeners = List(state, service, lifecycle, metrics, am)
+ val taskManager = new SamzaTaskManager(config, state, amClient, hConfig)
+
+ listeners = List(service, lifecycle, metrics, taskManager)
run(amClient, listeners, hConfig, interval)
} finally {
// jmxServer has to be stopped or will prevent process from exiting.
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterLifecycle.scala
----------------------------------------------------------------------
diff --git a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterLifecycle.scala b/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterLifecycle.scala
index d475c4c..2a5c0d8 100644
--- a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterLifecycle.scala
+++ b/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterLifecycle.scala
@@ -30,7 +30,7 @@ import org.apache.samza.util.Logging
* this means registering and unregistering with the RM, and shutting down
* when the RM tells us to Reboot.
*/
-class SamzaAppMasterLifecycle(containerMem: Int, containerCpu: Int, state: SamzaAppMasterState, amClient: AMRMClientAsync[ContainerRequest]) extends YarnAppMasterListener with Logging {
+class SamzaAppMasterLifecycle(containerMem: Int, containerCpu: Int, state: SamzaAppState, amClient: AMRMClientAsync[ContainerRequest]) extends YarnAppMasterListener with Logging {
var validResourceRequest = true
var shutdownMessage: String = null
@@ -50,7 +50,7 @@ class SamzaAppMasterLifecycle(containerMem: Int, containerCpu: Int, state: Samza
error(shutdownMessage)
validResourceRequest = false
state.status = FinalApplicationStatus.FAILED
- state.jobHealthy = false
+ state.jobHealthy.set(false)
}
}
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterMetrics.scala
----------------------------------------------------------------------
diff --git a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterMetrics.scala b/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterMetrics.scala
index 03acfe1..ae5a674 100644
--- a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterMetrics.scala
+++ b/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterMetrics.scala
@@ -46,7 +46,7 @@ object SamzaAppMasterMetrics {
*/
class SamzaAppMasterMetrics(
val config: Config,
- val state: SamzaAppMasterState,
+ val state: SamzaAppState,
val registry: ReadableMetricsRegistry) extends MetricsHelper with YarnAppMasterListener with Logging {
val jvm = new JvmMetrics(registry)
@@ -66,16 +66,16 @@ class SamzaAppMasterMetrics(
override def onInit() {
val mRunningContainers = newGauge("running-containers", () => state.runningContainers.size)
- val mNeededContainers = newGauge("needed-containers", () => state.neededContainers)
- val mCompletedContainers = newGauge("completed-containers", () => state.completedContainers)
- val mFailedContainers = newGauge("failed-containers", () => state.failedContainers)
- val mReleasedContainers = newGauge("released-containers", () => state.releasedContainers)
+ val mNeededContainers = newGauge("needed-containers", () => state.neededContainers.get())
+ val mCompletedContainers = newGauge("completed-containers", () => state.completedContainers.get())
+ val mFailedContainers = newGauge("failed-containers", () => state.failedContainers.get())
+ val mReleasedContainers = newGauge("released-containers", () => state.releasedContainers.get())
val mContainers = newGauge("container-count", () => state.containerCount)
val mHost = newGauge("http-host", () => state.nodeHost)
val mTrackingPort = newGauge("http-port", () => state.trackingUrl.getPort)
val mRpcPort = newGauge("rpc-port", () => state.rpcUrl.getPort)
val mAppAttemptId = newGauge("app-attempt-id", () => state.appAttemptId.toString)
- val mJobHealthy = newGauge("job-healthy", () => if (state.jobHealthy) 1 else 0)
+ val mJobHealthy = newGauge("job-healthy", () => if (state.jobHealthy.get()) 1 else 0)
jvm.start
reporters.values.foreach(_.start)
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterService.scala
----------------------------------------------------------------------
diff --git a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterService.scala b/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterService.scala
index 0605386..3adf86f 100644
--- a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterService.scala
+++ b/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterService.scala
@@ -35,7 +35,7 @@ import org.apache.samza.webapp.ApplicationMasterWebServlet
* dashboards to check on the status of a job. SamzaAppMasterService starts
* up the web service when initialized.
*/
-class SamzaAppMasterService(config: Config, state: SamzaAppMasterState, registry: ReadableMetricsRegistry, clientHelper: ClientHelper) extends YarnAppMasterListener with Logging {
+class SamzaAppMasterService(config: Config, state: SamzaAppState, registry: ReadableMetricsRegistry, clientHelper: ClientHelper) extends YarnAppMasterListener with Logging {
var rpcApp: HttpServer = null
var webApp: HttpServer = null
val SERVER_URL_OPT: String = "samza.autoscaling.server.url"
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterState.scala
----------------------------------------------------------------------
diff --git a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterState.scala b/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterState.scala
deleted file mode 100644
index f667c83..0000000
--- a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterState.scala
+++ /dev/null
@@ -1,56 +0,0 @@
-/*
- * 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.samza.job.yarn
-import org.apache.samza.util.Logging
-import org.apache.hadoop.yarn.api.records.FinalApplicationStatus
-import org.apache.hadoop.yarn.api.records.ContainerId
-import java.net.URL
-import org.apache.samza.coordinator.JobCoordinator
-
-/**
- * Samza's application master has state that is usually manipulated based on
- * responses from YARN's resource manager (via SamzaAppMasterTaskManager). This
- * class holds the current state of the application master.
- */
-class SamzaAppMasterState(val jobCoordinator: JobCoordinator, val taskId: Int, val containerId: ContainerId, val nodeHost: String, val nodePort: Int, val nodeHttpPort: Int) extends YarnAppMasterListener with Logging {
- // controlled by the AM
- var completedContainers = 0
- var neededContainers = 0
- var failedContainers = 0
- var releasedContainers = 0
- var containerCount = 0
- var runningContainers = Map[Int, YarnContainer]()
- var unclaimedContainers = Set[Int]()
- var finishedContainers = Set[Int]()
- var status = FinalApplicationStatus.UNDEFINED
- var jobHealthy = true
-
- // controlled by the service
- var trackingUrl: URL = null
- var rpcUrl: URL = null
- var coordinatorUrl: URL = null
-
- // controlled on startup
- var appAttemptId = containerId.getApplicationAttemptId
-
- // JMX address
- var jmxUrl = ""
- var jmxTunnelingUrl = ""
-}
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterTaskManager.scala
----------------------------------------------------------------------
diff --git a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterTaskManager.scala b/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterTaskManager.scala
deleted file mode 100644
index 1743c86..0000000
--- a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/SamzaAppMasterTaskManager.scala
+++ /dev/null
@@ -1,314 +0,0 @@
-/*
- * 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.samza.job.yarn
-
-import java.nio.ByteBuffer
-import java.util.Collections
-import scala.collection.JavaConversions._
-import org.apache.hadoop.fs.Path
-import org.apache.hadoop.io.DataOutputBuffer
-import org.apache.hadoop.security.UserGroupInformation
-import org.apache.hadoop.yarn.api.ApplicationConstants
-import org.apache.hadoop.yarn.api.protocolrecords.StartContainerRequest
-import org.apache.hadoop.yarn.api.records._
-import org.apache.hadoop.yarn.client.api.AMRMClient.ContainerRequest
-import org.apache.hadoop.yarn.client.api.NMClient
-import org.apache.hadoop.yarn.client.api.async.AMRMClientAsync
-import org.apache.hadoop.yarn.conf.YarnConfiguration
-import org.apache.hadoop.yarn.security.AMRMTokenIdentifier
-import org.apache.hadoop.yarn.util.ConverterUtils
-import org.apache.hadoop.yarn.util.Records
-import org.apache.samza.config.Config
-import org.apache.samza.config.TaskConfig.Config2Task
-import org.apache.samza.config.YarnConfig.Config2Yarn
-import org.apache.samza.job.CommandBuilder
-import org.apache.samza.job.ShellCommandBuilder
-import org.apache.samza.util.Logging
-import org.apache.samza.coordinator.JobCoordinator
-import org.apache.samza.util.{Util, Logging}
-import org.apache.samza.config.JobConfig.Config2Job
-
-object SamzaAppMasterTaskManager {
- val DEFAULT_CONTAINER_MEM = 1024
- val DEFAULT_CPU_CORES = 1
- val DEFAULT_CONTAINER_RETRY_COUNT = 8
- val DEFAULT_CONTAINER_RETRY_WINDOW_MS = 300000
-}
-
-case class ContainerFailure(val count: Int, val lastFailure: Long)
-
-/**
- * Samza's application master is mostly interested in requesting containers to
- * run Samza jobs. SamzaAppMasterTaskManager is responsible for requesting new
- * containers, handling failures, and notifying the application master that the
- * job is done.
- */
-class SamzaAppMasterTaskManager(clock: () => Long, config: Config, state: SamzaAppMasterState, amClient: AMRMClientAsync[ContainerRequest], conf: YarnConfiguration) extends YarnAppMasterListener with Logging {
- import SamzaAppMasterTaskManager._
-
- state.containerCount = config.getContainerCount
-
- var containerFailures = Map[Int, ContainerFailure]()
- var tooManyFailedContainers = false
- // TODO we might want to use NMClientAsync as well
- var containerManager: NMClient = null
-
- override def shouldShutdown = state.completedContainers == state.containerCount || tooManyFailedContainers
-
- override def onInit() {
- state.neededContainers = state.containerCount
- state.unclaimedContainers = state.jobCoordinator.jobModel.getContainers.keySet.map(_.toInt).toSet
- containerManager = NMClient.createNMClient()
- containerManager.init(conf)
- containerManager.start
-
- info("Requesting %s containers" format state.containerCount)
-
- requestContainers(config.getContainerMaxMemoryMb.getOrElse(DEFAULT_CONTAINER_MEM), config.getContainerMaxCpuCores.getOrElse(DEFAULT_CPU_CORES), state.neededContainers)
- }
-
- override def onShutdown {
- if (containerManager != null) {
- containerManager.stop
- }
- }
-
- override def onContainerAllocated(container: Container) {
- val containerIdStr = ConverterUtils.toString(container.getId)
-
- info("Got a container from YARN ResourceManager: %s" format container)
-
- state.unclaimedContainers.headOption match {
- case Some(containerId) => {
- info("Got available container ID (%d) for container: %s" format (containerId, container))
- val sspTaskNames = state.jobCoordinator.jobModel.getContainers.get(containerId)
- info("Claimed SSP taskNames %s for container ID %s" format (sspTaskNames, containerId))
- val cmdBuilderClassName = config.getCommandClass.getOrElse(classOf[ShellCommandBuilder].getName)
- val cmdBuilder = Class.forName(cmdBuilderClassName).newInstance.asInstanceOf[CommandBuilder]
- .setConfig(config)
- .setId(containerId)
- .setUrl(state.coordinatorUrl)
- val command = cmdBuilder.buildCommand
- info("Container ID %s using command %s" format (containerId, command))
- val env = cmdBuilder.buildEnvironment.map { case (k, v) => (k, Util.envVarEscape(v)) }
- info("Container ID %s using env %s" format (containerId, env))
- val path = new Path(config.getPackagePath.get)
- info("Starting container ID %s using package path %s" format (containerId, path))
-
- startContainer(
- path,
- container,
- env.toMap,
- "export SAMZA_LOG_DIR=%s && ln -sfn %s logs && exec ./__package/%s 1>logs/%s 2>logs/%s" format (ApplicationConstants.LOG_DIR_EXPANSION_VAR, ApplicationConstants.LOG_DIR_EXPANSION_VAR, command, ApplicationConstants.STDOUT, ApplicationConstants.STDERR))
-
- state.neededContainers -= 1
- if (state.neededContainers == 0) {
- state.jobHealthy = true
- }
- state.runningContainers += containerId -> new YarnContainer(container)
- state.unclaimedContainers -= containerId
-
- info("Claimed container ID %s for container %s on node %s (http://%s/node/containerlogs/%s)." format (containerId, containerIdStr, container.getNodeId.getHost, container.getNodeHttpAddress, containerIdStr))
-
- info("Started container ID %s" format containerId)
- }
- case _ => {
- // there are no more tasks to run, so release the container
- info("Got an extra container from YARN ResourceManager: %s" format (container))
-
- amClient.releaseAssignedContainer(container.getId)
- }
- }
- }
-
- override def onContainerCompleted(containerStatus: ContainerStatus) {
- val containerIdStr = ConverterUtils.toString(containerStatus.getContainerId)
- val containerId = state.runningContainers.filter { case (_, container) => container.id.equals(containerStatus.getContainerId()) }.keys.headOption
-
- containerId match {
- case Some(containerId) => {
- state.runningContainers -= containerId
- }
- case _ => None
- }
-
- containerStatus.getExitStatus match {
- case 0 => {
- info("Container %s completed successfully." format containerIdStr)
-
- state.completedContainers += 1
-
- if (containerId.isDefined) {
- state.finishedContainers += containerId.get
- containerFailures -= containerId.get
- }
-
- if (state.completedContainers == state.containerCount) {
- info("Setting job status to SUCCEEDED, since all containers have been marked as completed.")
- state.status = FinalApplicationStatus.SUCCEEDED
- }
- }
- case -100 => {
- info("Got an exit code of -100. This means that container %s was "
- + "killed by YARN, either due to being released by the application "
- + "master or being 'lost' due to node failures etc." format containerIdStr)
-
- state.releasedContainers += 1
-
- // If this container was assigned some partitions (a containerId), then
- // clean up, and request a new container for the tasks. This only
- // should happen if the container was 'lost' due to node failure, not
- // if the AM released the container.
- if (containerId.isDefined) {
- info("Released container %s was assigned task group ID %s. Requesting a new container for the task group." format (containerIdStr, containerId.get))
-
- state.neededContainers += 1
- state.jobHealthy = false
- state.unclaimedContainers += containerId.get
-
- // request a new container
- requestContainers(config.getContainerMaxMemoryMb.getOrElse(DEFAULT_CONTAINER_MEM), config.getContainerMaxCpuCores.getOrElse(DEFAULT_CPU_CORES), 1)
- }
- }
- case _ => {
- info("Container %s failed with exit code %d - %s." format (containerIdStr, containerStatus.getExitStatus, containerStatus.getDiagnostics))
-
- state.failedContainers += 1
- state.jobHealthy = false
-
- containerId match {
- case Some(containerId) =>
- info("Failed container %s owned task group ID %s." format (containerIdStr, containerId))
-
- state.unclaimedContainers += containerId
- state.neededContainers += 1
-
- // A container failed for an unknown reason. Let's check to see if
- // we need to shutdown the whole app master if too many container
- // failures have happened. The rules for failing are that the
- // failure count for a task group id must be > the configured retry
- // count, and the last failure (the one prior to this one) must have
- // happened less than retry window ms ago. If retry count is set to
- // 0, the app master will fail on any container failure. If the
- // retry count is set to a number < 0, a container failure will
- // never trigger an app master failure.
- val retryCount = config.getContainerRetryCount.getOrElse(DEFAULT_CONTAINER_RETRY_COUNT)
- val retryWindowMs = config.getContainerRetryWindowMs.getOrElse(DEFAULT_CONTAINER_RETRY_WINDOW_MS)
-
- if (retryCount == 0) {
- error("Container ID %s (%s) failed, and retry count is set to 0, so shutting down the application master, and marking the job as failed."
- format (containerId, containerIdStr))
-
- tooManyFailedContainers = true
- } else if (retryCount > 0) {
- val (currentFailCount, lastFailureTime) = containerFailures.get(containerId) match {
- case Some(ContainerFailure(count, lastFailure)) => (count + 1, lastFailure)
- case _ => (1, 0L)
- }
-
- if (currentFailCount > retryCount) {
- val lastFailureMsDiff = clock() - lastFailureTime
-
- if (lastFailureMsDiff < retryWindowMs) {
- error("Container ID %s (%s) has failed %s times, with last failure %sms ago. This is greater than retry count of %s and window of %s, so shutting down the application master, and marking the job as failed."
- format (containerId, containerIdStr, currentFailCount, lastFailureMsDiff, retryCount, retryWindowMs))
-
- // We have too many failures, and we're within the window
- // boundary, so reset shut down the app master.
- tooManyFailedContainers = true
- state.status = FinalApplicationStatus.FAILED
- } else {
- info("Resetting fail count for container ID %s back to 1, since last container failure (%s) for this container ID was outside the bounds of the retry window."
- format (containerId, containerIdStr))
-
- // Reset counter back to 1, since the last failure for this
- // container happened outside the window boundary.
- containerFailures += containerId -> ContainerFailure(1, clock())
- }
- } else {
- info("Current fail count for container ID %s is %s." format (containerId, currentFailCount))
- containerFailures += containerId -> ContainerFailure(currentFailCount, clock())
- }
- }
-
- if (!tooManyFailedContainers) {
- // Request a new container
- requestContainers(config.getContainerMaxMemoryMb.getOrElse(DEFAULT_CONTAINER_MEM), config.getContainerMaxCpuCores.getOrElse(DEFAULT_CPU_CORES), 1)
- }
- case _ => None
- }
- }
- }
- }
-
- protected def startContainer(packagePath: Path, container: Container, env: Map[String, String], cmds: String*) {
- info("starting container %s %s %s %s" format (packagePath, container, env, cmds))
-
- // set the local package so that the containers and app master are provisioned with it
- val packageResource = Records.newRecord(classOf[LocalResource])
- val packageUrl = ConverterUtils.getYarnUrlFromPath(packagePath)
- val fileStatus = packagePath.getFileSystem(conf).getFileStatus(packagePath)
-
- packageResource.setResource(packageUrl)
- packageResource.setSize(fileStatus.getLen)
- packageResource.setTimestamp(fileStatus.getModificationTime)
- packageResource.setType(LocalResourceType.ARCHIVE)
- packageResource.setVisibility(LocalResourceVisibility.APPLICATION)
-
- // copy tokens (copied from dist shell example)
- val credentials = UserGroupInformation.getCurrentUser.getCredentials
- val dob = new DataOutputBuffer
- credentials.writeTokenStorageToStream(dob)
- // now remove the AM->RM token so that containers cannot access it
- val iter = credentials.getAllTokens.iterator
- while (iter.hasNext) {
- val token = iter.next
- if (token.getKind.equals(AMRMTokenIdentifier.KIND_NAME)) {
- iter.remove
- }
- }
- val allTokens = ByteBuffer.wrap(dob.getData, 0, dob.getLength)
-
- // start the container
- val ctx = Records.newRecord(classOf[ContainerLaunchContext])
- ctx.setEnvironment(env)
- ctx.setTokens(allTokens.duplicate)
- ctx.setCommands(cmds.toList)
- ctx.setLocalResources(Collections.singletonMap("__package", packageResource))
-
- debug("setting package to %s" format packageResource)
- debug("setting context to %s" format ctx)
-
- val startContainerRequest = Records.newRecord(classOf[StartContainerRequest])
- startContainerRequest.setContainerLaunchContext(ctx)
- containerManager.startContainer(container, ctx)
- }
-
- protected def requestContainers(memMb: Int, cpuCores: Int, containers: Int) {
- info("Requesting %d container(s) with %dmb of memory" format (containers, memMb))
- val capability = Records.newRecord(classOf[Resource])
- val priority = Records.newRecord(classOf[Priority])
- priority.setPriority(0)
- capability.setMemory(memMb)
- capability.setVirtualCores(cpuCores)
- (0 until containers).foreach(idx => amClient.addContainerRequest(new ContainerRequest(capability, null, null, priority)))
- }
-
-}
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/main/scala/org/apache/samza/job/yarn/YarnJob.scala
----------------------------------------------------------------------
diff --git a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/YarnJob.scala b/samza-yarn/src/main/scala/org/apache/samza/job/yarn/YarnJob.scala
index 8dd70c9..02f46a1 100644
--- a/samza-yarn/src/main/scala/org/apache/samza/job/yarn/YarnJob.scala
+++ b/samza-yarn/src/main/scala/org/apache/samza/job/yarn/YarnJob.scala
@@ -31,7 +31,6 @@ import org.apache.samza.job.ApplicationStatus.Running
import org.apache.samza.job.StreamJob
import org.apache.samza.job.ApplicationStatus.SuccessfulFinish
import org.apache.samza.job.ApplicationStatus.UnsuccessfulFinish
-import org.apache.samza.config.YarnConfig.Config2Yarn
import org.apache.samza.config.JobConfig.Config2Job
import org.apache.samza.config.YarnConfig
import org.apache.samza.config.ShellCommandConfig
@@ -43,23 +42,20 @@ import org.apache.samza.config.MapConfig
import org.apache.samza.config.ConfigException
import org.apache.samza.config.SystemConfig
-object YarnJob {
- val DEFAULT_AM_CONTAINER_MEM = 1024
-}
/**
* Starts the application manager
*/
class YarnJob(config: Config, hadoopConfig: Configuration) extends StreamJob {
- import YarnJob._
val client = new ClientHelper(hadoopConfig)
var appId: Option[ApplicationId] = None
+ val yarnConfig = new YarnConfig(config)
def submit: YarnJob = {
appId = client.submitApplication(
- new Path(config.getPackagePath.getOrElse(throw new SamzaException("No YARN package path defined in config."))),
- config.getAMContainerMaxMemoryMb.getOrElse(DEFAULT_AM_CONTAINER_MEM),
+ new Path(yarnConfig.getPackagePath),
+ yarnConfig.getAMContainerMaxMemoryMb,
1,
List(
"export SAMZA_LOG_DIR=%s && ln -sfn %s logs && exec ./__package/bin/run-am.sh 1>logs/%s 2>logs/%s"
@@ -68,16 +64,17 @@ class YarnJob(config: Config, hadoopConfig: Configuration) extends StreamJob {
val coordinatorSystemConfig = Util.buildCoordinatorStreamConfig(config)
val envMap = Map(
ShellCommandConfig.ENV_COORDINATOR_SYSTEM_CONFIG -> Util.envVarEscape(SamzaObjectMapper.getObjectMapper.writeValueAsString(coordinatorSystemConfig)),
- ShellCommandConfig.ENV_JAVA_OPTS -> Util.envVarEscape(config.getAmOpts.getOrElse("")))
- val envMapWithJavaHome = config.getAMJavaHome match {
- case Some(javaHome) => envMap + (ShellCommandConfig.ENV_JAVA_HOME -> javaHome)
- case None => envMap
+ ShellCommandConfig.ENV_JAVA_OPTS -> Util.envVarEscape(yarnConfig.getAmOpts))
+ val amJavaHome = yarnConfig.getAMJavaHome
+ val envMapWithJavaHome = if(amJavaHome == null) {
+ envMap
+ } else {
+ envMap + (ShellCommandConfig.ENV_JAVA_HOME -> amJavaHome)
}
envMapWithJavaHome
}),
-
- Some("%s_%s" format (config.getName.get, config.getJobId.getOrElse(1))))
-
+ Some("%s_%s" format (config.getName.get, config.getJobId.getOrElse(1)))
+ )
this
}
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/main/scala/org/apache/samza/webapp/ApplicationMasterRestServlet.scala
----------------------------------------------------------------------
diff --git a/samza-yarn/src/main/scala/org/apache/samza/webapp/ApplicationMasterRestServlet.scala b/samza-yarn/src/main/scala/org/apache/samza/webapp/ApplicationMasterRestServlet.scala
index 09f4dc3..a40ab72 100644
--- a/samza-yarn/src/main/scala/org/apache/samza/webapp/ApplicationMasterRestServlet.scala
+++ b/samza-yarn/src/main/scala/org/apache/samza/webapp/ApplicationMasterRestServlet.scala
@@ -22,16 +22,14 @@ package org.apache.samza.webapp
import org.scalatra._
import scalate.ScalateSupport
import org.apache.samza.config.Config
-import org.apache.samza.job.yarn.SamzaAppMasterState
-import org.apache.samza.job.yarn.ClientHelper
+import org.apache.samza.job.yarn.{SamzaAppState, ClientHelper}
import org.apache.samza.metrics._
import scala.collection.JavaConversions._
import org.apache.hadoop.yarn.conf.YarnConfiguration
-import org.codehaus.jackson.map.ObjectMapper
import java.util.HashMap
import org.apache.samza.serializers.model.SamzaObjectMapper
-class ApplicationMasterRestServlet(config: Config, state: SamzaAppMasterState, registry: ReadableMetricsRegistry) extends ScalatraServlet with ScalateSupport {
+class ApplicationMasterRestServlet(config: Config, state: SamzaAppState, registry: ReadableMetricsRegistry) extends ScalatraServlet with ScalateSupport {
val yarnConfig = new YarnConfiguration
val client = new ClientHelper(yarnConfig)
val jsonMapper = SamzaObjectMapper.getObjectMapper
@@ -72,7 +70,7 @@ class ApplicationMasterRestServlet(config: Config, state: SamzaAppMasterState, r
val contextMap = new HashMap[String, Object]
contextMap.put("task-id", state.taskId: java.lang.Integer)
- contextMap.put("name", state.containerId.toString)
+ contextMap.put("name", state.amContainerId.toString)
jsonMapper.writeValueAsString(contextMap)
}
@@ -95,7 +93,7 @@ class ApplicationMasterRestServlet(config: Config, state: SamzaAppMasterState, r
val status = Map[String, Object](
"app-attempt-id" -> state.appAttemptId.toString,
- "container-id" -> state.containerId.toString,
+ "container-id" -> state.amContainerId.toString,
"containers" -> containers,
"host" -> "%s:%s".format(state.nodeHost, state.rpcUrl.getPort))
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/main/scala/org/apache/samza/webapp/ApplicationMasterWebServlet.scala
----------------------------------------------------------------------
diff --git a/samza-yarn/src/main/scala/org/apache/samza/webapp/ApplicationMasterWebServlet.scala b/samza-yarn/src/main/scala/org/apache/samza/webapp/ApplicationMasterWebServlet.scala
index 7fd5122..605332a 100644
--- a/samza-yarn/src/main/scala/org/apache/samza/webapp/ApplicationMasterWebServlet.scala
+++ b/samza-yarn/src/main/scala/org/apache/samza/webapp/ApplicationMasterWebServlet.scala
@@ -21,14 +21,14 @@ package org.apache.samza.webapp
import org.scalatra._
import scalate.ScalateSupport
-import org.apache.samza.job.yarn.SamzaAppMasterState
+import org.apache.samza.job.yarn.{SamzaAppState}
import org.apache.samza.config.Config
import scala.collection.JavaConversions._
import scala.collection.immutable.TreeMap
import org.apache.hadoop.yarn.conf.YarnConfiguration
import org.apache.hadoop.yarn.webapp.util.WebAppUtils
-class ApplicationMasterWebServlet(config: Config, state: SamzaAppMasterState) extends ScalatraServlet with ScalateSupport {
+class ApplicationMasterWebServlet(config: Config, state: SamzaAppState) extends ScalatraServlet with ScalateSupport {
val yarnConfig = new YarnConfiguration
before() {
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestContainerAllocator.java
----------------------------------------------------------------------
diff --git a/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestContainerAllocator.java b/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestContainerAllocator.java
new file mode 100644
index 0000000..0d07e28
--- /dev/null
+++ b/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestContainerAllocator.java
@@ -0,0 +1,245 @@
+/*
+ * 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.samza.job.yarn;
+
+import org.apache.hadoop.yarn.api.records.Container;
+import org.apache.hadoop.yarn.api.records.ContainerStatus;
+import org.apache.hadoop.yarn.client.api.async.impl.AMRMClientAsyncImpl;
+import org.apache.hadoop.yarn.util.ConverterUtils;
+import org.apache.samza.config.Config;
+import org.apache.samza.config.MapConfig;
+import org.apache.samza.container.TaskName;
+import org.apache.samza.coordinator.JobCoordinator;
+import org.apache.samza.coordinator.server.HttpServer;
+import org.apache.samza.job.model.ContainerModel;
+import org.apache.samza.job.model.JobModel;
+import org.apache.samza.job.model.TaskModel;
+import org.apache.samza.job.yarn.util.*;
+import org.apache.samza.job.yarn.util.TestUtil;
+import org.eclipse.jetty.servlet.DefaultServlet;
+import org.eclipse.jetty.servlet.ServletHolder;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.lang.reflect.Field;
+import java.net.URL;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.junit.Assert.*;
+
+public class TestContainerAllocator {
+ private final int ALLOCATOR_SLEEP_TIME = 10;
+ private static final String ANY_HOST = ContainerRequestState.ANY_HOST;
+ private final HttpServer server = new MockHttpServer("/", 7777, null, new ServletHolder(DefaultServlet.class));
+
+ private AMRMClientAsyncImpl amRmClientAsync;
+ private TestAMRMClientImpl testAMRMClient;
+ private ContainerAllocator containerAllocator;
+ private Thread allocatorThread;
+
+ private Config config = new MapConfig(new HashMap<String, String>() {
+ {
+ put("yarn.container.count", "1");
+ put("systems.test-system.samza.factory", "org.apache.samza.job.yarn.MockSystemFactory");
+ put("yarn.container.memory.mb", "512");
+ put("yarn.package.path", "/foo");
+ put("task.inputs", "test-system.test-stream");
+ put("systems.test-system.samza.key.serde", "org.apache.samza.serializers.JsonSerde");
+ put("systems.test-system.samza.msg.serde", "org.apache.samza.serializers.JsonSerde");
+ put("yarn.container.retry.count", "1");
+ put("yarn.container.retry.window.ms", "1999999999");
+ }
+ });
+
+ private SamzaAppState state = new SamzaAppState(getCoordinator(1), -1, ConverterUtils.toContainerId("container_1350670447861_0003_01_000001"), "", 1, 2);
+
+ private JobCoordinator getCoordinator(int containerCount) {
+ Map<Integer, ContainerModel> containers = new java.util.HashMap<>();
+ for (int i = 0; i < containerCount; i++) {
+ ContainerModel container = new ContainerModel(i, new HashMap<TaskName, TaskModel>());
+ containers.put(i, container);
+ }
+ JobModel jobModel = new JobModel(config, containers);
+ return new JobCoordinator(jobModel, server, null);
+ }
+
+ @Before
+ public void setup() throws Exception {
+ // Create AMRMClient
+ testAMRMClient = new TestAMRMClientImpl(
+ TestUtil.getAppMasterResponse(
+ false,
+ new ArrayList<Container>(),
+ new ArrayList<ContainerStatus>()
+ ));
+ amRmClientAsync = TestUtil.getAMClient(testAMRMClient);
+
+ // Initialize certain state variables (mostly to avoid NPE)
+ state.coordinatorUrl = new URL("http://localhost:7778/");
+
+ containerAllocator = new ContainerAllocator(
+ amRmClientAsync,
+ TestUtil.getContainerUtil(config, state),
+ ALLOCATOR_SLEEP_TIME
+ );
+ allocatorThread = new Thread(containerAllocator);
+ }
+
+ @After
+ public void teardown() throws Exception {
+ containerAllocator.setIsRunning(false);
+ allocatorThread.join();
+ }
+
+ /**
+ * Adds all containers returned to ANY_HOST only
+ */
+ @Test
+ public void testAddContainer() throws Exception {
+
+ Field requestStateField = containerAllocator.getClass().getSuperclass().getDeclaredField("containerRequestState");
+ requestStateField.setAccessible(true);
+
+ allocatorThread.start();
+
+ containerAllocator.requestContainers(new HashMap<Integer, String>() {
+ {
+ put(0, ANY_HOST);
+ put(1, ANY_HOST);
+ }
+ });
+
+ ContainerRequestState requestState = (ContainerRequestState) requestStateField.get(containerAllocator);
+
+ assertNull(requestState.getContainersOnAHost("abc"));
+ assertNull(requestState.getContainersOnAHost(ANY_HOST));
+
+ containerAllocator.addContainer(TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000002"), "abc", 123));
+ containerAllocator.addContainer(TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000003"), "xyz", 123));
+
+ assertNull(requestState.getContainersOnAHost("abc"));
+ assertNotNull(requestState.getContainersOnAHost(ANY_HOST));
+ assertTrue(requestState.getContainersOnAHost(ANY_HOST).size() == 2);
+ }
+
+ /**
+ * Test requestContainers
+ */
+ @Test
+ public void testRequestContainers() throws Exception {
+ Map<Integer, String> containersToHostMapping = new HashMap<Integer, String>() {
+ {
+ put(0, "abc");
+ put(1, "def");
+ put(2, null);
+ put(3, "abc");
+ }
+ };
+
+ allocatorThread.start();
+
+ containerAllocator.requestContainers(containersToHostMapping);
+
+ Field requestStateField = containerAllocator.getClass().getSuperclass().getDeclaredField("containerRequestState");
+ requestStateField.setAccessible(true);
+ ContainerRequestState requestState = (ContainerRequestState) requestStateField.get(containerAllocator);
+
+ assertNotNull(testAMRMClient.requests);
+ assertEquals(4, testAMRMClient.requests.size());
+
+ assertNotNull(requestState);
+
+ assertNotNull(requestState.getRequestsQueue());
+ assertTrue(requestState.getRequestsQueue().size() == 4);
+
+ // If host-affinty is not enabled, it doesn't update the requestMap
+ assertNotNull(requestState.getRequestsToCountMap());
+ assertTrue(requestState.getRequestsToCountMap().keySet().size() == 0);
+ }
+
+ /**
+ * Test request containers with no containerToHostMapping makes the right number of requests
+ */
+ @Test
+ public void testRequestContainersWithNoMapping() throws Exception {
+ int containerCount = 4;
+ Map<Integer, String> containersToHostMapping = new HashMap<Integer, String>();
+ for (int i = 0; i < containerCount; i++) {
+ containersToHostMapping.put(i, null);
+ }
+ allocatorThread.start();
+
+ containerAllocator.requestContainers(containersToHostMapping);
+
+ Field requestStateField = containerAllocator.getClass().getSuperclass().getDeclaredField("containerRequestState");
+ requestStateField.setAccessible(true);
+ ContainerRequestState requestState = (ContainerRequestState) requestStateField.get(containerAllocator);
+
+ assertNotNull(requestState);
+
+ assertNotNull(requestState.getRequestsQueue());
+ assertTrue(requestState.getRequestsQueue().size() == 4);
+
+ // If host-affinty is not enabled, it doesn't update the requestMap
+ assertNotNull(requestState.getRequestsToCountMap());
+ assertTrue(requestState.getRequestsToCountMap().keySet().size() == 0);
+ }
+
+ /**
+ * Extra allocated containers that are returned by the RM and unused by the AM should be released.
+ * Containers are considered "extra" only when there are no more pending requests to fulfill
+ * @throws Exception
+ */
+ @Test
+ public void testAllocatorReleasesExtraContainers() throws Exception {
+ Container container = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000001"), "abc", 123);
+ Container container1 = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000002"), "abc", 123);
+ Container container2 = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000003"), "def", 123);
+
+ allocatorThread.start();
+
+ containerAllocator.requestContainer(0, "abc");
+
+ containerAllocator.addContainer(container);
+ containerAllocator.addContainer(container1);
+ containerAllocator.addContainer(container2);
+
+ Thread.sleep(600);
+
+ Field requestStateField = containerAllocator.getClass().getSuperclass().getDeclaredField("containerRequestState");
+ requestStateField.setAccessible(true);
+ ContainerRequestState requestState = (ContainerRequestState) requestStateField.get(containerAllocator);
+
+ assertNotNull(testAMRMClient.getRelease());
+ assertEquals(2, testAMRMClient.getRelease().size());
+ assertTrue(testAMRMClient.getRelease().contains(container1.getId()));
+ assertTrue(testAMRMClient.getRelease().contains(container2.getId()));
+
+ // Test that state is cleaned up
+ assertEquals(0, requestState.getRequestsQueue().size());
+ assertEquals(0, requestState.getRequestsToCountMap().size());
+ assertNull(requestState.getContainersOnAHost("abc"));
+ assertNull(requestState.getContainersOnAHost("def"));
+ }
+
+}
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestContainerRequestState.java
----------------------------------------------------------------------
diff --git a/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestContainerRequestState.java b/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestContainerRequestState.java
new file mode 100644
index 0000000..402fe78
--- /dev/null
+++ b/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestContainerRequestState.java
@@ -0,0 +1,221 @@
+/*
+ * 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.samza.job.yarn;
+
+import org.apache.hadoop.yarn.api.records.Container;
+import org.apache.hadoop.yarn.api.records.ContainerStatus;
+import org.apache.hadoop.yarn.client.api.AMRMClient;
+import org.apache.hadoop.yarn.client.api.async.impl.AMRMClientAsyncImpl;
+import org.apache.hadoop.yarn.util.ConverterUtils;
+import org.apache.samza.job.yarn.util.TestAMRMClientImpl;
+import org.apache.samza.job.yarn.util.TestUtil;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.ArrayList;
+
+import static org.junit.Assert.*;
+
+public class TestContainerRequestState {
+ private AMRMClientAsyncImpl amRmClientAsync;
+ private TestAMRMClientImpl testAMRMClient;
+ private static final String ANY_HOST = ContainerRequestState.ANY_HOST;
+
+ @Before
+ public void setup() {
+ // Create AMRMClient
+ testAMRMClient = new TestAMRMClientImpl(
+ TestUtil.getAppMasterResponse(
+ false,
+ new ArrayList<Container>(),
+ new ArrayList<ContainerStatus>()
+ ));
+ amRmClientAsync = TestUtil.getAMClient(testAMRMClient);
+ }
+
+ /**
+ * Test state after a request is submitted
+ */
+ @Test
+ public void testUpdateRequestState() {
+ // Host-affinity is enabled
+ ContainerRequestState state = new ContainerRequestState(amRmClientAsync, true);
+ SamzaContainerRequest request = new SamzaContainerRequest(0, "abc");
+ state.updateRequestState(request);
+
+ assertNotNull(testAMRMClient.requests);
+ assertEquals(1, testAMRMClient.requests.size());
+ assertEquals(request.getIssuedRequest(), testAMRMClient.requests.get(0));
+
+ assertNotNull(state.getRequestsQueue());
+ assertTrue(state.getRequestsQueue().size() == 1);
+
+ assertNotNull(state.getRequestsToCountMap());
+ assertNotNull(state.getRequestsToCountMap().get("abc"));
+ assertEquals(1, state.getRequestsToCountMap().get("abc").get());
+
+ assertNotNull(state.getContainersOnAHost("abc"));
+ assertEquals(0, state.getContainersOnAHost("abc").size());
+
+ // Host-affinity is not enabled
+ ContainerRequestState state1 = new ContainerRequestState(amRmClientAsync, false);
+ SamzaContainerRequest request1 = new SamzaContainerRequest(1, null);
+ state1.updateRequestState(request1);
+
+ assertNotNull(testAMRMClient.requests);
+ assertEquals(2, testAMRMClient.requests.size());
+
+ AMRMClient.ContainerRequest expectedContainerRequest = request1.getIssuedRequest();
+ AMRMClient.ContainerRequest actualContainerRequest = testAMRMClient.requests.get(1);
+
+ assertEquals(expectedContainerRequest.getCapability(), actualContainerRequest.getCapability());
+ assertEquals(expectedContainerRequest.getPriority(), actualContainerRequest.getPriority());
+ assertNull(actualContainerRequest.getNodes());
+
+ assertNotNull(state1.getRequestsQueue());
+ assertTrue(state1.getRequestsQueue().size() == 1);
+
+ assertNotNull(state1.getRequestsToCountMap());
+ assertNull(state1.getRequestsToCountMap().get(ANY_HOST));
+
+ }
+
+ /**
+ * Test addContainer() updates the state correctly
+ */
+ @Test
+ public void testAddContainer() {
+ // Add container to ANY_LIST when host-affinity is not enabled
+ ContainerRequestState state = new ContainerRequestState(amRmClientAsync, false);
+ Container container = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000002"), "abc", 123);
+ state.addContainer(container);
+
+ assertNotNull(state.getRequestsQueue());
+ assertNotNull(state.getRequestsToCountMap());
+ assertNotNull(state.getContainersOnAHost(ANY_HOST));
+
+ assertEquals(1, state.getContainersOnAHost(ANY_HOST).size());
+ assertEquals(container, state.getContainersOnAHost(ANY_HOST).get(0));
+
+ // Container Allocated when there is no request in queue
+ ContainerRequestState state1 = new ContainerRequestState(amRmClientAsync, true);
+ Container container1 = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000003"), "zzz", 123);
+ state1.addContainer(container1);
+
+ assertNotNull(state1.getRequestsQueue());
+ assertEquals(0, state1.getRequestsQueue().size());
+
+ assertNull(state1.getContainersOnAHost("zzz"));
+ assertNotNull(state1.getContainersOnAHost(ANY_HOST));
+ assertEquals(1, state1.getContainersOnAHost(ANY_HOST).size());
+ assertEquals(container1, state1.getContainersOnAHost(ANY_HOST).get(0));
+
+ // Container Allocated on a Requested Host
+ state1.updateRequestState(new SamzaContainerRequest(0, "abc"));
+
+ assertNotNull(state1.getRequestsQueue());
+ assertEquals(1, state1.getRequestsQueue().size());
+
+ assertNotNull(state1.getRequestsToCountMap());
+ assertNotNull(state1.getRequestsToCountMap().get("abc"));
+ assertEquals(1, state1.getRequestsToCountMap().get("abc").get());
+
+ state1.addContainer(container);
+
+ assertNotNull(state1.getContainersOnAHost("abc"));
+ assertEquals(1, state1.getContainersOnAHost("abc").size());
+ assertEquals(container, state1.getContainersOnAHost("abc").get(0));
+
+ // Container Allocated on host that was not requested
+ Container container2 = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000004"), "xyz", 123);
+ state1.addContainer(container2);
+
+ assertNull(state1.getContainersOnAHost("xyz"));
+ assertNotNull(state1.getContainersOnAHost(ANY_HOST));
+ assertEquals(2, state1.getContainersOnAHost(ANY_HOST).size());
+ assertEquals(container2, state1.getContainersOnAHost(ANY_HOST).get(1));
+
+ // Extra containers were allocated on a host that was requested
+ Container container3 = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000005"), "abc", 123);
+ state1.addContainer(container3);
+
+ assertEquals(3, state1.getContainersOnAHost(ANY_HOST).size());
+ assertEquals(container3, state1.getContainersOnAHost(ANY_HOST).get(2));
+ }
+
+ /**
+ * Test request state after container is assigned to a host
+ * * Assigned on requested host
+ * * Assigned on any host
+ */
+ @Test
+ public void testContainerAssignment() throws Exception {
+ // Host-affinity enabled
+ ContainerRequestState state = new ContainerRequestState(amRmClientAsync, true);
+ SamzaContainerRequest request = new SamzaContainerRequest(0, "abc");
+ SamzaContainerRequest request1 = new SamzaContainerRequest(0, "def");
+
+ state.updateRequestState(request);
+ state.updateRequestState(request1);
+
+ Container container = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000001"), "abc", 123);
+ Container container1 = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000001"), "zzz", 123);
+ state.addContainer(container);
+ state.addContainer(container1);
+
+ assertEquals(2, state.getRequestsQueue().size());
+ assertEquals(2, state.getRequestsToCountMap().size());
+
+ assertNotNull(state.getContainersOnAHost("abc"));
+ assertEquals(1, state.getContainersOnAHost("abc").size());
+ assertEquals(container, state.getContainersOnAHost("abc").get(0));
+
+ assertNotNull(state.getContainersOnAHost("def"));
+ assertEquals(0, state.getContainersOnAHost("def").size());
+
+ assertNotNull(state.getContainersOnAHost(ANY_HOST));
+ assertEquals(1, state.getContainersOnAHost(ANY_HOST).size());
+ assertEquals(container1, state.getContainersOnAHost(ANY_HOST).get(0));
+
+ // Container assigned on the requested host
+ state.updateStateAfterAssignment(request, "abc", container);
+
+ assertEquals(1, state.getRequestsQueue().size());
+ assertEquals(request1, state.getRequestsQueue().peek());
+
+ assertNotNull(state.getRequestsToCountMap().get("abc"));
+ assertEquals(0, state.getRequestsToCountMap().get("abc").get());
+
+ assertNotNull(state.getContainersOnAHost("abc"));
+ assertEquals(0, state.getContainersOnAHost("abc").size());
+
+ // Container assigned on any host
+ state.updateStateAfterAssignment(request1, ANY_HOST, container1);
+
+ assertEquals(0, state.getRequestsQueue().size());
+
+ assertNotNull(state.getRequestsToCountMap().get("def"));
+ assertEquals(0, state.getRequestsToCountMap().get("def").get());
+
+ assertNotNull(state.getContainersOnAHost(ANY_HOST));
+ assertEquals(0, state.getContainersOnAHost(ANY_HOST).size());
+
+ }
+
+}
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestHostAwareContainerAllocator.java
----------------------------------------------------------------------
diff --git a/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestHostAwareContainerAllocator.java b/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestHostAwareContainerAllocator.java
new file mode 100644
index 0000000..ee4a5b0
--- /dev/null
+++ b/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestHostAwareContainerAllocator.java
@@ -0,0 +1,335 @@
+/*
+ * 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.samza.job.yarn;
+
+import org.apache.hadoop.yarn.api.records.Container;
+import org.apache.hadoop.yarn.api.records.ContainerStatus;
+import org.apache.hadoop.yarn.client.api.async.impl.AMRMClientAsyncImpl;
+import org.apache.hadoop.yarn.util.ConverterUtils;
+import org.apache.samza.config.Config;
+import org.apache.samza.config.MapConfig;
+import org.apache.samza.container.TaskName;
+import org.apache.samza.coordinator.JobCoordinator;
+import org.apache.samza.coordinator.server.HttpServer;
+import org.apache.samza.job.model.ContainerModel;
+import org.apache.samza.job.model.JobModel;
+import org.apache.samza.job.model.TaskModel;
+import org.apache.samza.job.yarn.util.MockContainerUtil;
+import org.apache.samza.job.yarn.util.MockHttpServer;
+import org.apache.samza.job.yarn.util.TestAMRMClientImpl;
+import org.apache.samza.job.yarn.util.TestUtil;
+import org.eclipse.jetty.servlet.DefaultServlet;
+import org.eclipse.jetty.servlet.ServletHolder;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.lang.reflect.Field;
+import java.net.URL;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+public class TestHostAwareContainerAllocator {
+ private static final int ALLOCATOR_SLEEP_TIME = 1;
+ private static final int CONTAINER_REQUEST_TIMEOUT = 3;
+ private static final String ANY_HOST = ContainerRequestState.ANY_HOST;
+
+ private final HttpServer server = new MockHttpServer("/", 7777, null, new ServletHolder(DefaultServlet.class));
+
+ private AMRMClientAsyncImpl amRmClientAsync;
+ private TestAMRMClientImpl testAMRMClient;
+ private HostAwareContainerAllocator containerAllocator;
+ private Thread allocatorThread;
+ private ContainerUtil containerUtil;
+
+ private Config config = new MapConfig(new HashMap<String, String>() {
+ {
+ put("yarn.container.count", "1");
+ put("systems.test-system.samza.factory", "org.apache.samza.job.yarn.MockSystemFactory");
+ put("yarn.container.memory.mb", "512");
+ put("yarn.package.path", "/foo");
+ put("task.inputs", "test-system.test-stream");
+ put("systems.test-system.samza.key.serde", "org.apache.samza.serializers.JsonSerde");
+ put("systems.test-system.samza.msg.serde", "org.apache.samza.serializers.JsonSerde");
+ put("yarn.container.retry.count", "1");
+ put("yarn.container.retry.window.ms", "1999999999");
+ put("yarn.samza.host-affinity.enabled", "true");
+ }
+ });
+
+ private Config getConfig() {
+ Map<String, String> map = new HashMap<>();
+ map.putAll(config);
+ return new MapConfig(map);
+ }
+
+ private SamzaAppState state = new SamzaAppState(getCoordinator(1), -1, ConverterUtils.toContainerId("container_1350670447861_0003_01_000001"), "", 1, 2);
+
+ private JobCoordinator getCoordinator(int containerCount) {
+ Map<Integer, ContainerModel> containers = new java.util.HashMap<>();
+ for (int i = 0; i < containerCount; i++) {
+ ContainerModel container = new ContainerModel(i, new HashMap<TaskName, TaskModel>());
+ containers.put(i, container);
+ }
+ JobModel jobModel = new JobModel(getConfig(), containers);
+ return new JobCoordinator(jobModel, server, null);
+ }
+
+
+ @Before
+ public void setup() throws Exception {
+ // Create AMRMClient
+ testAMRMClient = new TestAMRMClientImpl(
+ TestUtil.getAppMasterResponse(
+ false,
+ new ArrayList<Container>(),
+ new ArrayList<ContainerStatus>()
+ ));
+ amRmClientAsync = TestUtil.getAMClient(testAMRMClient);
+
+ // Initialize certain state variables
+ state.coordinatorUrl = new URL("http://localhost:7778/");
+
+ containerUtil = TestUtil.getContainerUtil(getConfig(), state);
+
+ containerAllocator = new HostAwareContainerAllocator(
+ amRmClientAsync,
+ containerUtil,
+ ALLOCATOR_SLEEP_TIME,
+ CONTAINER_REQUEST_TIMEOUT
+ );
+ allocatorThread = new Thread(containerAllocator);
+ }
+
+ @After
+ public void teardown() throws Exception {
+ containerAllocator.setIsRunning(false);
+ allocatorThread.join();
+ }
+
+ @Test
+ public void testAllocatorReleasesExtraContainers() throws Exception {
+ Container container = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000001"), "abc", 123);
+ Container container1 = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000002"), "abc", 123);
+ Container container2 = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000003"), "def", 123);
+
+ allocatorThread.start();
+
+ containerAllocator.requestContainer(0, "abc");
+
+ containerAllocator.addContainer(container);
+ containerAllocator.addContainer(container1);
+ containerAllocator.addContainer(container2);
+
+ Thread.sleep(600);
+
+ Field requestStateField = containerAllocator.getClass().getSuperclass().getDeclaredField("containerRequestState");
+ requestStateField.setAccessible(true);
+ ContainerRequestState requestState = (ContainerRequestState) requestStateField.get(containerAllocator);
+
+ assertNotNull(testAMRMClient.getRelease());
+ assertEquals(2, testAMRMClient.getRelease().size());
+ assertTrue(testAMRMClient.getRelease().contains(container1.getId()));
+ assertTrue(testAMRMClient.getRelease().contains(container2.getId()));
+
+ // Test that state is cleaned up
+ assertEquals(0, requestState.getRequestsQueue().size());
+ assertEquals(0, requestState.getRequestsToCountMap().size());
+ assertNull(requestState.getContainersOnAHost("abc"));
+ assertNull(requestState.getContainersOnAHost("def"));
+ }
+
+ /**
+ * Test request containers with no containerToHostMapping makes the right number of requests
+ */
+ @Test
+ public void testRequestContainersWithNoMapping() throws Exception {
+ int containerCount = 4;
+ Map<Integer, String> containersToHostMapping = new HashMap<Integer, String>();
+ for (int i = 0; i < containerCount; i++) {
+ containersToHostMapping.put(i, null);
+ }
+
+ allocatorThread.start();
+
+ containerAllocator.requestContainers(containersToHostMapping);
+
+ Field requestStateField = containerAllocator.getClass().getSuperclass().getDeclaredField("containerRequestState");
+ requestStateField.setAccessible(true);
+ ContainerRequestState requestState = (ContainerRequestState) requestStateField.get(containerAllocator);
+
+ assertNotNull(requestState);
+
+ assertNotNull(requestState.getRequestsQueue());
+ assertEquals(4, requestState.getRequestsQueue().size());
+
+ assertNotNull(requestState.getRequestsToCountMap());
+ assertEquals(1, requestState.getRequestsToCountMap().keySet().size());
+ assertTrue(requestState.getRequestsToCountMap().keySet().contains(ANY_HOST));
+ }
+
+ /**
+ * Add containers to the correct host in the request state
+ */
+ @Test
+ public void testAddContainerWithHostAffinity() throws Exception {
+ Field requestStateField = containerAllocator.getClass().getSuperclass().getDeclaredField("containerRequestState");
+ requestStateField.setAccessible(true);
+ ContainerRequestState requestState = (ContainerRequestState) requestStateField.get(containerAllocator);
+
+ containerAllocator.requestContainers(new HashMap<Integer, String>() {
+ {
+ put(0, "abc");
+ put(1, "xyz");
+ }
+ });
+
+ allocatorThread.start();
+
+ assertNotNull(requestState.getContainersOnAHost("abc"));
+ assertEquals(0, requestState.getContainersOnAHost("abc").size());
+
+ assertNotNull(requestState.getContainersOnAHost("xyz"));
+ assertEquals(0, requestState.getContainersOnAHost("xyz").size());
+
+ assertNull(requestState.getContainersOnAHost(ANY_HOST));
+
+ containerAllocator.addContainer(TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000002"), "abc", 123));
+ containerAllocator.addContainer(TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000004"), "def", 123));
+ containerAllocator.addContainer(TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000003"), "xyz", 123));
+
+ assertNotNull(requestState.getContainersOnAHost("abc"));
+ assertTrue(requestState.getContainersOnAHost("abc").size() == 1);
+
+ assertNotNull(requestState.getContainersOnAHost("xyz"));
+ assertTrue(requestState.getContainersOnAHost("xyz").size() == 1);
+
+ assertNotNull(requestState.getContainersOnAHost(ANY_HOST));
+ assertTrue(requestState.getContainersOnAHost(ANY_HOST).size() == 1);
+ assertEquals(ConverterUtils.toContainerId("container_1350670447861_0003_01_000004"), requestState.getContainersOnAHost(ANY_HOST).get(0).getId());
+ }
+
+ @Test
+ public void testRequestContainers() throws Exception {
+ Map<Integer, String> containersToHostMapping = new HashMap<Integer, String>() {
+ {
+ put(0, "abc");
+ put(1, "def");
+ put(2, null);
+ put(3, "abc");
+ }
+ };
+ allocatorThread.start();
+
+ containerAllocator.requestContainers(containersToHostMapping);
+
+
+ Field requestStateField = containerAllocator.getClass().getSuperclass().getDeclaredField("containerRequestState");
+ requestStateField.setAccessible(true);
+ ContainerRequestState requestState = (ContainerRequestState) requestStateField.get(containerAllocator);
+
+ assertNotNull(testAMRMClient.requests);
+ assertEquals(4, testAMRMClient.requests.size());
+
+ assertNotNull(requestState);
+
+ assertNotNull(requestState.getRequestsQueue());
+ assertTrue(requestState.getRequestsQueue().size() == 4);
+
+ assertNotNull(requestState.getRequestsToCountMap());
+ Map<String, AtomicInteger> requestsMap = requestState.getRequestsToCountMap();
+
+ assertNotNull(requestsMap.get("abc"));
+ assertEquals(2, requestsMap.get("abc").get());
+
+ assertNotNull(requestsMap.get("def"));
+ assertEquals(1, requestsMap.get("def").get());
+
+ assertNotNull(requestsMap.get(ANY_HOST));
+ assertEquals(1, requestsMap.get(ANY_HOST).get());
+ }
+
+ /**
+ * Handles expired requests correctly and assigns ANY_HOST
+ */
+ @Test
+ public void testExpiredRequestHandling() throws Exception {
+ Map<Integer, String> containersToHostMapping = new HashMap<Integer, String>() {
+ {
+ put(0, "abc");
+ put(1, "def");
+ }
+ };
+ containerAllocator.requestContainers(containersToHostMapping);
+
+ Field requestStateField = containerAllocator.getClass().getSuperclass().getDeclaredField("containerRequestState");
+ requestStateField.setAccessible(true);
+ ContainerRequestState requestState = (ContainerRequestState) requestStateField.get(containerAllocator);
+
+ assertNotNull(requestState.getRequestsQueue());
+ assertTrue(requestState.getRequestsQueue().size() == 2);
+
+ assertNotNull(requestState.getRequestsToCountMap());
+ assertNotNull(requestState.getRequestsToCountMap().get("abc"));
+ assertTrue(requestState.getRequestsToCountMap().get("abc").get() == 1);
+
+ assertNotNull(requestState.getRequestsToCountMap().get("def"));
+ assertTrue(requestState.getRequestsToCountMap().get("def").get() == 1);
+
+ allocatorThread.start();
+
+ Thread.sleep(600);
+
+ Container container0 = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000002"), "xyz", 123);
+ Container container1 = TestUtil.getContainer(ConverterUtils.toContainerId("container_1350670447861_0003_01_000003"), "zzz", 123);
+ containerAllocator.addContainer(container0);
+ containerAllocator.addContainer(container1);
+
+ assertNull(requestState.getContainersOnAHost("xyz"));
+ assertNull(requestState.getContainersOnAHost("zzz"));
+ assertNotNull(requestState.getContainersOnAHost(ANY_HOST));
+ assertTrue(requestState.getContainersOnAHost(ANY_HOST).size() == 2);
+
+ Thread.sleep(1000);
+
+ MockContainerUtil mockContainerUtil = (MockContainerUtil) containerUtil;
+
+ assertNotNull(mockContainerUtil.runningContainerList.get("xyz"));
+ assertTrue(mockContainerUtil.runningContainerList.get("xyz").contains(container0));
+
+ assertNotNull(mockContainerUtil.runningContainerList.get("zzz"));
+ assertTrue(mockContainerUtil.runningContainerList.get("zzz").contains(container1));
+
+ assertNull(requestState.getContainersOnAHost(ANY_HOST));
+
+ assertNotNull(requestState.getRequestsQueue());
+ assertTrue(requestState.getRequestsQueue().size() == 0);
+
+ assertNotNull(requestState.getRequestsToCountMap());
+ assertNull(requestState.getRequestsToCountMap().get("abc"));
+
+ assertNull(requestState.getRequestsToCountMap().get("def"));
+ }
+}
http://git-wip-us.apache.org/repos/asf/samza/blob/c76132ff/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestSamzaContainerRequest.java
----------------------------------------------------------------------
diff --git a/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestSamzaContainerRequest.java b/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestSamzaContainerRequest.java
new file mode 100644
index 0000000..aefae65
--- /dev/null
+++ b/samza-yarn/src/test/java/org/apache/samza/job/yarn/TestSamzaContainerRequest.java
@@ -0,0 +1,43 @@
+/*
+ * 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.samza.job.yarn;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+public class TestSamzaContainerRequest {
+ private static final String ANY_HOST = ContainerRequestState.ANY_HOST;
+
+ @Test
+ public void testPreferredHostIsNeverNull() {
+ SamzaContainerRequest request = new SamzaContainerRequest(0, null);
+
+ assertNotNull(request.getPreferredHost());
+
+ // preferredHost is null, it should automatically default to ANY_HOST
+ assertTrue(request.getPreferredHost().equals(ANY_HOST));
+
+ SamzaContainerRequest request1 = new SamzaContainerRequest(1, "abc");
+ assertNotNull(request1.getPreferredHost());
+ assertTrue(request1.getPreferredHost().equals("abc"));
+
+ }
+}