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

Fixes #14292.

### Description

When `spark.rapids.memory.host.offHeapLimit.enabled` is false (the
default),
`getPinnedPoolAndOffHeapLimits` returns `-1` for the non-pinned limit
(meaning unlimited).
The code that sets the default MT read memory limit was using this `-1`
literally:

```scala
(0.9 * (pinnedSize + nonPinnedLimit)).toLong  // nonPinnedLimit == -1
```

With no pinned pool this produces `0`; with a pinned pool (e.g. 2 GiB)
it produces `~0.9 *
pinnedSize`, sizing the limit off the pinned pool only and ignoring all
non-pinned memory.

The fix extracts the hardware-derived off-heap limit computation from
`getPinnedPoolAndOffHeapLimits`
into a shared `computeEffectiveOffHeapLimit` method, and uses it in
`computeMtReadLimit` when
`nonPinnedLimit == -1` to derive a concrete total from hardware instead.

Note: `spark.rapids.sql.multiThreadedRead.memoryLimit.enabled` defaults
to `false`, so impact is
limited to users who have explicitly enabled the feature. The
hardware-derived calculation has a
known approximation limitation in certain environments (see #13628),
called out in a new code
comment.

**Before (2 GiB pinned pool, off-heap limit disabled):** MT read limit ≈
1.80 GiB (90% of pinned only)

**After:** MT read limit ≈ 78 GiB (90% of hardware-derived off-heap)

### Checklists

- [ ] This PR has added documentation for new or modified features or
behaviors.
- [x] This PR has added new tests or modified existing tests to cover
new code paths.
`GpuDeviceManagerSuite` — "MT read limit should be 90% of total off heap
when off heap limit
disabled": calls `computeMtReadLimit` with `nonPinnedLimit=-1` and
asserts the result equals
90% of the hardware-derived off-heap total. Fails against pre-fix code,
passes after.
- [ ] Performance testing has been performed and its results are added
in the PR description. Or, an issue has been filed with a link in the PR
description.

---------

Signed-off-by: Zach Puller <zpuller@nvidia.com>

Author: zpuller

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

@@ -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
-    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
     // queries to complete, whereas this limit was sufficient.
@@ -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) {

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

@@ -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.
@@ -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"