client+server:use mem group size in transport

Signed-off-by: zihao zhao <[email protected]>
NVIDIA · Jun 3, 2024 · 99b8009 · 99b8009
1 parent 5abb350
commit 99b8009
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)
     }
   }