client+server:use mem group size in transport

NVIDIA · Jun 3, 2024 · b846bb4 · b846bb4
1 parent 5abb350
commit b846bb4
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/main/scala/org/apache/spark/shuffle/ucx/external/ExternalUcxTransport.scala b/src/main/scala/org/apache/spark/shuffle/ucx/external/ExternalUcxTransport.scala
@@ -63,7 +63,8 @@ class ExternalUcxTransport(val ucxShuffleConf: ExternalUcxConf) extends UcxLoggi
                        ucxShuffleConf.minRegistrationSize,
                        ucxShuffleConf.maxRegistrationSize / numPools,
                        ucxShuffleConf.preallocateBuffersMap,
-                       ucxShuffleConf.memoryLimit)
+                       ucxShuffleConf.memoryLimit,
+                       ucxShuffleConf.memoryGroupSize)
     }
   }