Bring back setRapidsShuffleManager in the driver side (#3389)

Signed-off-by: Alessandro Bellina <[email protected]>
NVIDIA · Sep 4, 2021 · 7cb995c · 7cb995c
1 parent 83b58d3
commit 7cb995c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/RapidsShuffleInternalManagerBase.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/RapidsShuffleInternalManagerBase.scala
@@ -236,6 +236,8 @@ abstract class RapidsShuffleInternalManagerBase(conf: SparkConf, val isDriver: B
 
   protected val wrapped = new SortShuffleManager(conf)
 
+  GpuShuffleEnv.setRapidsShuffleManager(Some(this))
+
   private[this] val transportEnabledMessage = if (!rapidsConf.shuffleTransportEnabled) {
     "Transport disabled (local cached blocks only)"
   } else {