Skip to content

Commit f0ec9a2

Browse files
zpullerzpuller
committed
base mt read limits off effective off heap limits
Signed-off-by: Zach Puller <zpuller@nvidia.com>
1 parent edde7e1 commit f0ec9a2

2 files changed

Lines changed: 128 additions & 67 deletions

File tree

‎sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuDeviceManager.scala‎

Lines changed: 102 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -485,22 +485,38 @@ object GpuDeviceManager extends Logging {
485485
// if it is not set already.
486486
if (conf.multiThreadReadMemoryLimit == 0) {
487487
sparkConf.set(RapidsConf.MULTITHREAD_READ_MEMORY_LIMIT_SIZE.key,
488-
(0.9 * (pinnedSize + nonPinnedLimit)).toLong.toString)
488+
computeMtReadLimit(pinnedSize, nonPinnedLimit, conf, sparkConf,
489+
Cuda.getDeviceCount).toString)
489490
}
490491
}
491492

492493
// visible for testing
493-
def getPinnedPoolAndOffHeapLimits(conf: RapidsConf, sparkConf: SparkConf, deviceCount: Int,
494-
memCheck: MemoryChecker =
495-
MemoryCheckerImpl): (Long, Long) = {
496-
val perTaskOverhead = conf.perTaskOverhead
497-
val totalOverhead = perTaskOverhead * GpuDeviceManager.numCores
498-
val confPinnedSize = conf.pinnedPoolSize
494+
def computeMtReadLimit(pinnedSize: Long, nonPinnedLimit: Long,
495+
conf: RapidsConf, sparkConf: SparkConf, deviceCount: Int,
496+
memCheck: MemoryChecker = MemoryCheckerImpl): Long = {
497+
val totalOffHeap = if (nonPinnedLimit >= 0) {
498+
pinnedSize + nonPinnedLimit
499+
} else {
500+
// nonPinnedLimit == -1 means off-heap limit tracking is disabled (unlimited).
501+
// Derive a concrete total from hardware so the MT read limit is still meaningful.
502+
// Note: this feature is off by default
503+
// (spark.rapids.sql.multiThreadedRead.memoryLimit.enabled defaults to false), so the risk
504+
// of this approximation affecting users is low. The
505+
// underlying hardware-derived calculation has known limitations — in particular, it may not
506+
// account correctly for spark.memory.offHeap.size in all environments (see
507+
// https://github.com/NVIDIA/spark-rapids/issues/13628). For now we reuse the same
508+
// best-effort estimate rather than special-casing this path.
509+
computeEffectiveOffHeapLimit(conf, sparkConf, deviceCount, memCheck)
510+
}
511+
val limit = (0.9 * totalOffHeap).toLong
512+
logInfo(s"Setting MT read memory limit to ${limit / 1024 / 1024.0} MiB " +
513+
s"(90% of ${totalOffHeap / 1024 / 1024.0} MiB total off-heap)")
514+
limit
515+
}
516+
517+
private def computeEffectiveOffHeapLimit(conf: RapidsConf, sparkConf: SparkConf,
518+
deviceCount: Int, memCheck: MemoryChecker): Long = {
499519
val confLimit = conf.offHeapLimit
500-
val confLimitEnabled = conf.offHeapLimitEnabled
501-
// This min limit of 4GB is somewhat arbitrary, but based on some testing which showed
502-
// that the previous minimum of 15 MB * num cores was too little for certain benchmark
503-
// queries to complete, whereas this limit was sufficient.
504520
val minMemoryLimit = 4L * 1024 * 1024 * 1024
505521

506522
val executorOverheadKey = "spark.executor.memoryOverhead"
@@ -521,68 +537,87 @@ object GpuDeviceManager extends Logging {
521537
0L
522538
}
523539

524-
if (confLimitEnabled) {
525-
val memoryLimit = if (confLimit.isDefined) {
526-
if (executorOverhead.isEmpty) {
527-
logWarning(s"$executorOverheadKey is not set")
528-
}
529-
logInfo(s"using configured ${RapidsConf.OFF_HEAP_LIMIT_SIZE} of ${confLimit.get}")
530-
confLimit.get
540+
if (confLimit.isDefined) {
541+
if (executorOverhead.isEmpty) {
542+
logWarning(s"$executorOverheadKey is not set")
543+
}
544+
logInfo(s"using configured ${RapidsConf.OFF_HEAP_LIMIT_SIZE} of ${confLimit.get}")
545+
confLimit.get
546+
} else {
547+
// in case we cannot query the host for available memory due to environmental
548+
// constraints, getAvailableMemoryBytes returns None and we fall back to 0, which causes
549+
// basedOnHostMemory to go negative and trip the minMemoryLimit floor below.
550+
lazy val availableHostMemory = if (isIntegratedGpu) {
551+
(memCheck.getAvailableMemoryBytes(conf).getOrElse(0L) * (1.0 -
552+
conf.integratedGpuMemoryFraction)).toLong
531553
} else {
532-
// in case we cannot query the host for available memory due to environmental
533-
// constraints, we can fall back to minMemoryLimit via saying there's no available
534-
lazy val availableHostMemory = if (isIntegratedGpu) {
535-
(memCheck.getAvailableMemoryBytes(conf).getOrElse(0L) * (1.0 -
536-
conf.integratedGpuMemoryFraction)).toLong
554+
memCheck.getAvailableMemoryBytes(conf).getOrElse(0L)
555+
}
556+
val hostMemUsageFraction = .8
557+
// Spark calculates the total mem to allocate to the job as
558+
// val totalMemMiB =
559+
// executorMemoryMiB + memoryOverheadMiB + memoryOffHeapMiB + pysparkMemToUseMiB
560+
// and RAPIDS uses memory from the overhead portion here. Therefore, if the overhead is
561+
// set we can just use that, otherwise we can infer it from the above as
562+
// val memoryOverheadMiB =
563+
// totalMemMiB - executorMemoryMiB - memoryOffHeapMiB - pysparkMemToUseMiB
564+
// where totalMemMiB is instead derived from the actual mem limits we can observe
565+
// directly from the system.
566+
// Note: subtracting sparkOffHeapSize here is an approximation. In some environments
567+
// (e.g. Databricks) spark.memory.offHeap.size is set to a large value and RAPIDS is
568+
// expected to operate within that allocation rather than outside it, so subtracting it
569+
// double-counts the memory and under-estimates the available limit. Fixing this properly
570+
// may require hooking into Spark's MemoryManager. See
571+
// https://github.com/NVIDIA/spark-rapids/issues/13628
572+
lazy val basedOnHostMemory = (hostMemUsageFraction * ((1.0 * availableHostMemory /
573+
deviceCount) - heapSize - pysparkOverhead - sparkOffHeapSize)).toLong
574+
if (executorOverhead.isDefined) {
575+
val basedOnConfiguredOverhead = executorOverhead.get
576+
logWarning(s"${RapidsConf.OFF_HEAP_LIMIT_SIZE} is not set; we derived " +
577+
s"a memory limit from ($executorOverheadKey = ${executorOverhead.get}")
578+
if (basedOnConfiguredOverhead < minMemoryLimit) {
579+
logWarning(s"memory limit $basedOnConfiguredOverhead is less than the minimum of " +
580+
s"$minMemoryLimit; using the latter")
581+
if (minMemoryLimit > basedOnHostMemory) {
582+
logWarning(s"the amount of available memory detected on the host is " +
583+
s"$availableHostMemory, based off of which we computed a limit of " +
584+
s"$basedOnHostMemory, which is less than the minimum $minMemoryLimit, " +
585+
s"so we are using the minimum $minMemoryLimit")
586+
}
587+
minMemoryLimit
537588
} else {
538-
memCheck.getAvailableMemoryBytes(conf).getOrElse(0L)
589+
basedOnConfiguredOverhead
539590
}
540-
val hostMemUsageFraction = .8
541-
// Spark calculates the total mem to allocate to the job as
542-
// val totalMemMiB =
543-
// executorMemoryMiB + memoryOverheadMiB + memoryOffHeapMiB + pysparkMemToUseMiB
544-
// and RAPIDS uses memory from the overhead portion here. Therefore, if the overhead is
545-
// set we can just use that, otherwise we can infer it from the above as
546-
// val memoryOverheadMiB =
547-
// totalMemMiB - executorMemoryMiB - memoryOffHeapMiB - pysparkMemToUseMiB
548-
// where totalMemMiB is instead derived from the actual mem limits we can observe
549-
// directly from the system
550-
lazy val basedOnHostMemory = (hostMemUsageFraction * ((1.0 * availableHostMemory /
551-
deviceCount) - heapSize - pysparkOverhead - sparkOffHeapSize)).toLong
552-
if (executorOverhead.isDefined) {
553-
val basedOnConfiguredOverhead = executorOverhead.get
554-
logWarning(s"${RapidsConf.OFF_HEAP_LIMIT_SIZE} is not set; we derived " +
555-
s"a memory limit from ($executorOverheadKey = ${executorOverhead.get}")
556-
if (basedOnConfiguredOverhead < minMemoryLimit) {
557-
logWarning(s"memory limit $basedOnConfiguredOverhead is less than the minimum of " +
558-
s"$minMemoryLimit; using the latter")
559-
if (minMemoryLimit > basedOnHostMemory) {
560-
logWarning(s"the amount of available memory detected on the host is " +
561-
s"$availableHostMemory, based off of which we computed a limit of " +
562-
s"$basedOnHostMemory, which is less than the minimum $minMemoryLimit, " +
563-
s"so we are using the minimum $minMemoryLimit")
564-
}
565-
minMemoryLimit
566-
} else {
567-
basedOnConfiguredOverhead
568-
}
591+
} else {
592+
logWarning(s"${RapidsConf.OFF_HEAP_LIMIT_SIZE} is not set; we used " +
593+
s"memory limit derived from ($hostMemUsageFraction * (estimated available " +
594+
s"host memory / device count) - $heapSizeKey - $pysparkOverheadKey - " +
595+
s"$sparkOffHeapSizeKey) = ($hostMemUsageFraction * ($availableHostMemory / " +
596+
s"$deviceCount) - $heapSize - $pysparkOverhead - $sparkOffHeapSize) = " +
597+
s"$basedOnHostMemory")
598+
if (basedOnHostMemory < minMemoryLimit) {
599+
logWarning(s"the memory limit, $basedOnHostMemory, based on the available " +
600+
s"host memory of $availableHostMemory, is less than the minimum of " +
601+
s"$minMemoryLimit; using the latter $minMemoryLimit")
602+
minMemoryLimit
569603
} else {
570-
logWarning(s"${RapidsConf.OFF_HEAP_LIMIT_SIZE} is not set; we used " +
571-
s"memory limit derived from ($hostMemUsageFraction * (estimated available " +
572-
s"host memory / device count) - $heapSizeKey - $pysparkOverheadKey - " +
573-
s"$sparkOffHeapSizeKey) = ($hostMemUsageFraction * ($availableHostMemory / " +
574-
s"$deviceCount) - $heapSize - $pysparkOverhead - $sparkOffHeapSize) = " +
575-
s"$basedOnHostMemory")
576-
if (basedOnHostMemory < minMemoryLimit) {
577-
logWarning(s"the memory limit, $basedOnHostMemory, based on the available " +
578-
s"host memory of $availableHostMemory, is less than the minimum of " +
579-
s"$minMemoryLimit; using the latter $minMemoryLimit")
580-
minMemoryLimit
581-
} else {
582-
basedOnHostMemory
583-
}
604+
basedOnHostMemory
584605
}
585606
}
607+
}
608+
}
609+
610+
// visible for testing
611+
def getPinnedPoolAndOffHeapLimits(conf: RapidsConf, sparkConf: SparkConf, deviceCount: Int,
612+
memCheck: MemoryChecker =
613+
MemoryCheckerImpl): (Long, Long) = {
614+
val perTaskOverhead = conf.perTaskOverhead
615+
val totalOverhead = perTaskOverhead * GpuDeviceManager.numCores
616+
val confPinnedSize = conf.pinnedPoolSize
617+
val confLimitEnabled = conf.offHeapLimitEnabled
618+
619+
if (confLimitEnabled) {
620+
val memoryLimit = computeEffectiveOffHeapLimit(conf, sparkConf, deviceCount, memCheck)
586621

587622
// Now we need to know the pinned vs non-pinned limits
588623
val pinnedLimit = if (confPinnedSize + totalOverhead <= memoryLimit) {

‎tests/src/test/scala/com/nvidia/spark/rapids/GpuDeviceManagerSuite.scala‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -437,6 +437,32 @@ class GpuDeviceManagerSuite extends AnyFunSuite with BeforeAndAfter {
437437
}
438438
}
439439

440+
test("MT read limit should be 90% of total off heap when off heap limit disabled") {
441+
val deviceCount = 1
442+
val availableHostMem = 32L * 1024 * 1024 * 1024 // 32 GiB
443+
val heapSize = toBytes("1g") // default spark.executor.memory
444+
445+
val sparkConf = new SparkConf()
446+
val rapidsConf = new RapidsConf(Map(
447+
RapidsConf.OFF_HEAP_LIMIT_ENABLED.key -> "false"))
448+
449+
TestMemoryChecker.setAvailableMemoryBytes(Some(availableHostMem))
450+
451+
try {
452+
val mtLimit = GpuDeviceManager.computeMtReadLimit(
453+
pinnedSize = 0L, nonPinnedLimit = -1L,
454+
rapidsConf, sparkConf, deviceCount, TestMemoryChecker)
455+
456+
// When off-heap limit is disabled, the effective total off-heap should be derived
457+
// from hardware (same 80% formula used in the enabled path), and MT read limit = 90% of that.
458+
val effectiveOffHeap = (0.8 * (availableHostMem - heapSize)).toLong
459+
val expectedMtLimit = (0.9 * effectiveOffHeap).toLong
460+
assertResult(expectedMtLimit)(mtLimit)
461+
} finally {
462+
TestMemoryChecker.setAvailableMemoryBytes(None)
463+
}
464+
}
465+
440466
test("get host memory limits with discrete GPU") {
441467
val deviceCount = 1
442468
val pySparkOverheadStr = "2g"

0 commit comments

Comments
 (0)