diff --git 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
index abd37834ed3cc..4c2d7075f736a 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
@@ -31,6 +31,7 @@ import org.apache.hadoop.yarn.api.AMRMProtocol
 import org.apache.hadoop.yarn.api.records._
 import org.apache.hadoop.yarn.api.protocolrecords.AllocateRequest
 import org.apache.hadoop.yarn.util.Records
+import org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse
 
 /**
  * Acquires resources for executors from a ResourceManager and launches executors in new containers.
@@ -118,7 +119,7 @@ private[yarn] class YarnAllocationHandler(
           request.getPriority,
           request.getCapability))
     }
-    new AlphaAllocateResponse(resourceManager.allocate(req).getAMResponse())
+    new AlphaAllocateResponse(resourceManager.allocate(req))
   }
 
   override protected def releaseContainer(container: Container) = {
@@ -220,7 +221,7 @@ private[yarn] class YarnAllocationHandler(
     retval
   }
 
-  private class AlphaAllocateResponse(response: AMResponse) extends YarnAllocateResponse {
+  private class AlphaAllocateResponse(response: AllocateResponse) extends YarnAllocateResponse {
     override def getAllocatedContainers() = response.getAllocatedContainers()
     override def getAvailableResources() = response.getAvailableResources()
     override def getCompletedContainersStatuses() = response.getCompletedContainersStatuses()