Skip to content

Fix MT read memory limit defaulting to wrong size when off-heap limit is disabled #14507

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zpuller merged 2 commits into from
Apr 7, 2026
Merged

Fix MT read memory limit defaulting to wrong size when off-heap limit is disabled #14507

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f0ec9a2
base mt read limits off effective off heap limits
zpuller Mar 31, 2026
c0eeeb8
pr comments
zpuller Apr 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
171 changes: 105 additions & 66 deletions sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuDeviceManager.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -481,23 +481,43 @@ object GpuDeviceManager extends Logging {
}
// Host memory limits must be set after the pinned memory pool is initialized
HostAlloc.initialize(nonPinnedLimit)
// Fill the MULTITHREAD_READ_MEMORY_LIMIT_SIZE with the 90% of the total OFF_HEAP memory
// if it is not set already.
// Fill the MULTITHREAD_READ_MEMORY_LIMIT_SIZE with 90% of the total off-heap memory
// if it is not set already. When off-heap limit tracking is disabled (nonPinnedLimit == -1),
// falls back to a hardware-derived estimate via computeEffectiveOffHeapLimit.
if (conf.multiThreadReadMemoryLimit == 0) {
sparkConf.set(RapidsConf.MULTITHREAD_READ_MEMORY_LIMIT_SIZE.key,
(0.9 * (pinnedSize + nonPinnedLimit)).toLong.toString)
computeMtReadLimit(pinnedSize, nonPinnedLimit, conf, sparkConf,
Cuda.getDeviceCount).toString)
}
}

// visible for testing
def getPinnedPoolAndOffHeapLimits(conf: RapidsConf, sparkConf: SparkConf, deviceCount: Int,
memCheck: MemoryChecker =
MemoryCheckerImpl): (Long, Long) = {
val perTaskOverhead = conf.perTaskOverhead
val totalOverhead = perTaskOverhead * GpuDeviceManager.numCores
val confPinnedSize = conf.pinnedPoolSize
def computeMtReadLimit(pinnedSize: Long, nonPinnedLimit: Long,
conf: RapidsConf, sparkConf: SparkConf, deviceCount: Int,
memCheck: MemoryChecker = MemoryCheckerImpl): Long = {
val totalOffHeap = if (nonPinnedLimit >= 0) {
pinnedSize + nonPinnedLimit
} else {
// nonPinnedLimit == -1 means off-heap limit tracking is disabled (unlimited).
// Derive a concrete total from hardware so the MT read limit is still meaningful.
// Note: this feature is off by default
// (spark.rapids.sql.multiThreadedRead.memoryLimit.enabled defaults to false), so the risk
// of this approximation affecting users is low. The
// underlying hardware-derived calculation has known limitations — in particular, it may not
// account correctly for spark.memory.offHeap.size in all environments (see
// https://github.com/NVIDIA/spark-rapids/issues/13628). For now we reuse the same
// best-effort estimate rather than special-casing this path.
computeEffectiveOffHeapLimit(conf, sparkConf, deviceCount, memCheck)
}
val limit = (0.9 * totalOffHeap).toLong
logInfo(s"Setting MT read memory limit to ${limit / 1024 / 1024.0} MiB " +
s"(90% of ${totalOffHeap / 1024 / 1024.0} MiB total off-heap)")
limit
}

private def computeEffectiveOffHeapLimit(conf: RapidsConf, sparkConf: SparkConf,
deviceCount: Int, memCheck: MemoryChecker): Long = {
val confLimit = conf.offHeapLimit
Comment thread
zpuller marked this conversation as resolved.
val confLimitEnabled = conf.offHeapLimitEnabled
// This min limit of 4GB is somewhat arbitrary, but based on some testing which showed
// that the previous minimum of 15 MB * num cores was too little for certain benchmark
Comment thread
zpuller marked this conversation as resolved.
// queries to complete, whereas this limit was sufficient.
Expand All @@ -521,68 +541,87 @@ object GpuDeviceManager extends Logging {
0L
}

if (confLimitEnabled) {
val memoryLimit = if (confLimit.isDefined) {
if (executorOverhead.isEmpty) {
logWarning(s"$executorOverheadKey is not set")
}
logInfo(s"using configured ${RapidsConf.OFF_HEAP_LIMIT_SIZE} of ${confLimit.get}")
confLimit.get
if (confLimit.isDefined) {
if (executorOverhead.isEmpty) {
logWarning(s"$executorOverheadKey is not set")
}
logInfo(s"using configured ${RapidsConf.OFF_HEAP_LIMIT_SIZE} of ${confLimit.get}")
confLimit.get
} else {
// in case we cannot query the host for available memory due to environmental
// constraints, getAvailableMemoryBytes returns None and we fall back to 0, which causes
// basedOnHostMemory to go negative and trip the minMemoryLimit floor below.
lazy val availableHostMemory = if (isIntegratedGpu) {
(memCheck.getAvailableMemoryBytes(conf).getOrElse(0L) * (1.0 -
conf.integratedGpuMemoryFraction)).toLong
} else {
// in case we cannot query the host for available memory due to environmental
// constraints, we can fall back to minMemoryLimit via saying there's no available
lazy val availableHostMemory = if (isIntegratedGpu) {
(memCheck.getAvailableMemoryBytes(conf).getOrElse(0L) * (1.0 -
conf.integratedGpuMemoryFraction)).toLong
memCheck.getAvailableMemoryBytes(conf).getOrElse(0L)
}
val hostMemUsageFraction = .8
// Spark calculates the total mem to allocate to the job as
// val totalMemMiB =
// executorMemoryMiB + memoryOverheadMiB + memoryOffHeapMiB + pysparkMemToUseMiB
// and RAPIDS uses memory from the overhead portion here. Therefore, if the overhead is
// set we can just use that, otherwise we can infer it from the above as
// val memoryOverheadMiB =
// totalMemMiB - executorMemoryMiB - memoryOffHeapMiB - pysparkMemToUseMiB
// where totalMemMiB is instead derived from the actual mem limits we can observe
// directly from the system.
// Note: subtracting sparkOffHeapSize here is an approximation. In some environments
// (e.g. Databricks) spark.memory.offHeap.size is set to a large value and RAPIDS is
// expected to operate within that allocation rather than outside it, so subtracting it
// double-counts the memory and under-estimates the available limit. Fixing this properly
// may require hooking into Spark's MemoryManager. See
// https://github.com/NVIDIA/spark-rapids/issues/13628
lazy val basedOnHostMemory = (hostMemUsageFraction * ((1.0 * availableHostMemory /
deviceCount) - heapSize - pysparkOverhead - sparkOffHeapSize)).toLong
if (executorOverhead.isDefined) {
val basedOnConfiguredOverhead = executorOverhead.get
logWarning(s"${RapidsConf.OFF_HEAP_LIMIT_SIZE} is not set; we derived " +
s"a memory limit from ($executorOverheadKey = ${executorOverhead.get}")
if (basedOnConfiguredOverhead < minMemoryLimit) {
logWarning(s"memory limit $basedOnConfiguredOverhead is less than the minimum of " +
s"$minMemoryLimit; using the latter")
if (minMemoryLimit > basedOnHostMemory) {
logWarning(s"the amount of available memory detected on the host is " +
s"$availableHostMemory, based off of which we computed a limit of " +
s"$basedOnHostMemory, which is less than the minimum $minMemoryLimit, " +
s"so we are using the minimum $minMemoryLimit")
}
minMemoryLimit
} else {
memCheck.getAvailableMemoryBytes(conf).getOrElse(0L)
basedOnConfiguredOverhead
}
val hostMemUsageFraction = .8
// Spark calculates the total mem to allocate to the job as
// val totalMemMiB =
// executorMemoryMiB + memoryOverheadMiB + memoryOffHeapMiB + pysparkMemToUseMiB
// and RAPIDS uses memory from the overhead portion here. Therefore, if the overhead is
// set we can just use that, otherwise we can infer it from the above as
// val memoryOverheadMiB =
// totalMemMiB - executorMemoryMiB - memoryOffHeapMiB - pysparkMemToUseMiB
// where totalMemMiB is instead derived from the actual mem limits we can observe
// directly from the system
lazy val basedOnHostMemory = (hostMemUsageFraction * ((1.0 * availableHostMemory /
deviceCount) - heapSize - pysparkOverhead - sparkOffHeapSize)).toLong
if (executorOverhead.isDefined) {
val basedOnConfiguredOverhead = executorOverhead.get
logWarning(s"${RapidsConf.OFF_HEAP_LIMIT_SIZE} is not set; we derived " +
s"a memory limit from ($executorOverheadKey = ${executorOverhead.get}")
if (basedOnConfiguredOverhead < minMemoryLimit) {
logWarning(s"memory limit $basedOnConfiguredOverhead is less than the minimum of " +
s"$minMemoryLimit; using the latter")
if (minMemoryLimit > basedOnHostMemory) {
logWarning(s"the amount of available memory detected on the host is " +
s"$availableHostMemory, based off of which we computed a limit of " +
s"$basedOnHostMemory, which is less than the minimum $minMemoryLimit, " +
s"so we are using the minimum $minMemoryLimit")
}
minMemoryLimit
} else {
basedOnConfiguredOverhead
}
} else {
logWarning(s"${RapidsConf.OFF_HEAP_LIMIT_SIZE} is not set; we used " +
s"memory limit derived from ($hostMemUsageFraction * (estimated available " +
s"host memory / device count) - $heapSizeKey - $pysparkOverheadKey - " +
s"$sparkOffHeapSizeKey) = ($hostMemUsageFraction * ($availableHostMemory / " +
s"$deviceCount) - $heapSize - $pysparkOverhead - $sparkOffHeapSize) = " +
s"$basedOnHostMemory")
if (basedOnHostMemory < minMemoryLimit) {
logWarning(s"the memory limit, $basedOnHostMemory, based on the available " +
s"host memory of $availableHostMemory, is less than the minimum of " +
s"$minMemoryLimit; using the latter $minMemoryLimit")
minMemoryLimit
} else {
logWarning(s"${RapidsConf.OFF_HEAP_LIMIT_SIZE} is not set; we used " +
s"memory limit derived from ($hostMemUsageFraction * (estimated available " +
s"host memory / device count) - $heapSizeKey - $pysparkOverheadKey - " +
s"$sparkOffHeapSizeKey) = ($hostMemUsageFraction * ($availableHostMemory / " +
s"$deviceCount) - $heapSize - $pysparkOverhead - $sparkOffHeapSize) = " +
s"$basedOnHostMemory")
if (basedOnHostMemory < minMemoryLimit) {
logWarning(s"the memory limit, $basedOnHostMemory, based on the available " +
s"host memory of $availableHostMemory, is less than the minimum of " +
s"$minMemoryLimit; using the latter $minMemoryLimit")
minMemoryLimit
} else {
basedOnHostMemory
}
basedOnHostMemory
}
}
}
}

// visible for testing
def getPinnedPoolAndOffHeapLimits(conf: RapidsConf, sparkConf: SparkConf, deviceCount: Int,
memCheck: MemoryChecker =
MemoryCheckerImpl): (Long, Long) = {
val perTaskOverhead = conf.perTaskOverhead
val totalOverhead = perTaskOverhead * GpuDeviceManager.numCores
val confPinnedSize = conf.pinnedPoolSize
val confLimitEnabled = conf.offHeapLimitEnabled

if (confLimitEnabled) {
val memoryLimit = computeEffectiveOffHeapLimit(conf, sparkConf, deviceCount, memCheck)

// Now we need to know the pinned vs non-pinned limits
val pinnedLimit = if (confPinnedSize + totalOverhead <= memoryLimit) {
Expand Down
28 changes: 27 additions & 1 deletion tests/src/test/scala/com/nvidia/spark/rapids/GpuDeviceManagerSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2020-2025, NVIDIA CORPORATION.
* Copyright (c) 2020-2026, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -437,6 +437,32 @@ class GpuDeviceManagerSuite extends AnyFunSuite with BeforeAndAfter {
}
}

test("MT read limit should be 90% of total off heap when off heap limit disabled") {
val deviceCount = 1
val availableHostMem = 32L * 1024 * 1024 * 1024 // 32 GiB
val heapSize = toBytes("1g") // default spark.executor.memory

val sparkConf = new SparkConf()
val rapidsConf = new RapidsConf(Map(
RapidsConf.OFF_HEAP_LIMIT_ENABLED.key -> "false"))

TestMemoryChecker.setAvailableMemoryBytes(Some(availableHostMem))

try {
val mtLimit = GpuDeviceManager.computeMtReadLimit(
pinnedSize = 0L, nonPinnedLimit = -1L,
rapidsConf, sparkConf, deviceCount, TestMemoryChecker)

// When off-heap limit is disabled, the effective total off-heap should be derived
// from hardware (same 80% formula used in the enabled path), and MT read limit = 90% of that.
val effectiveOffHeap = (0.8 * (availableHostMem - heapSize)).toLong
val expectedMtLimit = (0.9 * effectiveOffHeap).toLong
assertResult(expectedMtLimit)(mtLimit)
} finally {
TestMemoryChecker.setAvailableMemoryBytes(None)
}
}

test("get host memory limits with discrete GPU") {
val deviceCount = 1
val pySparkOverheadStr = "2g"
Expand Down