Count backward eps in numEps hint. (#23)

Signed-off-by: Peter Rudenko <[email protected]> Signed-off-by: Peter Rudenko <[email protected]>
NVIDIA · Sep 22, 2022 · 493d6b0 · 493d6b0
1 parent fb18bc5
commit 493d6b0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/main/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransport.scala b/src/main/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransport.scala
@@ -106,7 +106,8 @@ class UcxShuffleTransport(var ucxShuffleConf: UcxShuffleConf = null, var executo
  }
 
  val numEndpoints = ucxShuffleConf.numWorkers *
- ucxShuffleConf.getSparkConf.getInt("spark.executor.instances", 1)
+ ucxShuffleConf.getSparkConf.getInt("spark.executor.instances", 1) *
+ ucxShuffleConf.numListenerThreads // Each listener thread creates backward endpoint
  logInfo(s"Creating UCX context with an estimated number of endpoints: $numEndpoints")
 
  val params = new UcpParams().requestAmFeature().setMtWorkersShared(true).setEstimatedNumEps(numEndpoints)