Use the shuffle service port specified through hadoop config

Author: Andrew Or

core/src/main/scala/org/apache/spark/storage/BlockManager.scala

@@ -30,6 +30,7 @@ import akka.actor.{ActorSystem, Props}
 import sun.nio.ch.DirectBuffer
 
 import org.apache.spark._
+import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.executor._
 import org.apache.spark.io.CompressionCodec
 import org.apache.spark.network._
@@ -92,7 +93,19 @@ private[spark] class BlockManager(
 
   private[spark]
   val externalShuffleServiceEnabled = conf.getBoolean("spark.shuffle.service.enabled", false)
-  private val externalShuffleServicePort = conf.getInt("spark.shuffle.service.port", 7337)
+
+  // In Yarn, the shuffle service port maybe set through the Hadoop config
+  private val shuffleServicePortKey = "spark.shuffle.service.port"
+  private val externalShuffleServicePort = {
+    val sparkPort = conf.getInt(shuffleServicePortKey, 7337)
+    if (SparkHadoopUtil.get.isYarnMode) {
+      val hadoopConf = SparkHadoopUtil.get.newConfiguration(conf)
+      Option(hadoopConf.get(shuffleServicePortKey)).map(_.toInt).getOrElse(sparkPort)
+    } else {
+      sparkPort
+    }
+  }
+
   // Check that we're not using external shuffle service with consolidated shuffle files.
   if (externalShuffleServiceEnabled
       && conf.getBoolean("spark.shuffle.consolidateFiles", false)

network/yarn/src/main/java/org/apache/spark/network/yarn/YarnShuffleService.java

@@ -49,6 +49,7 @@ protected void serviceInit(Configuration conf) {
       RpcHandler rpcHandler = new ExternalShuffleBlockHandler();
       TransportContext transportContext = new TransportContext(transportConf, rpcHandler);
       transportContext.createServer(port);
+      logger.info("Started Yarn shuffle service for Spark on port " + port);
     } catch (Exception e) {
       logger.error("Exception in starting Yarn shuffle service for Spark", e);
     }