Count backward eps in numEps hint. (#23)

petro-rudenko · web-flow · commit 493d6b02f8f7 · 2022-09-22T12:52:42.000+03:00
Signed-off-by: Peter Rudenko &lt;prudenko@nvidia.com&gt;

Signed-off-by: Peter Rudenko &lt;prudenko@nvidia.com&gt;
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)

Original file line number	Diff line number	Diff line change
`@@ -106,7 +106,8 @@ class UcxShuffleTransport(var ucxShuffleConf: UcxShuffleConf = null, var executo`
`106`	`106`	`}`
`107`	`107`
`108`	`108`	`val numEndpoints = ucxShuffleConf.numWorkers *`
`109`		`- ucxShuffleConf.getSparkConf.getInt("spark.executor.instances", 1)`
	`109`	`+ ucxShuffleConf.getSparkConf.getInt("spark.executor.instances", 1) *`
	`110`	`+ ucxShuffleConf.numListenerThreads // Each listener thread creates backward endpoint`
`110`	`111`	`logInfo(s"Creating UCX context with an estimated number of endpoints: $numEndpoints")`
`111`	`112`
`112`	`113`	`val params = new UcpParams().requestAmFeature().setMtWorkersShared(true).setEstimatedNumEps(numEndpoints)`