From 3de8dcedb70bb10fcb144be30bec3b11583a9a68 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Fri, 24 Jul 2026 16:55:14 +0800 Subject: [PATCH] [GLUTEN-12621][CORE] Fix NoSuchElementException in untracked memory mode without off-heap size setPredefinedConfigs read the off-heap size in its non-dynamic branch with the single-arg conf.getSizeAsBytes(SPARK_OFFHEAP_SIZE_KEY). checkOffHeapSettings returns early for untracked memory mode and skips the requirement that the size be set, so the key can be absent there. SparkConf.get(String) throws NoSuchElementException on a missing key, so driver plugin init crashed instead of proceeding. VeloxListenerApi.onDriverStart has the same single-arg read and runs after setPredefinedConfigs, so the failure reaches end to end. The dynamic-sizing branch already shields it by setting SPARK_OFFHEAP_SIZE_KEY to "0"; untracked mode did not. Normalize the key to 0 in the untracked branch when it is absent, mirroring the dynamic branch, so one change covers both readers. Normal mode is unaffected since checkOffHeapSettings guarantees the key is set. Add GlutenDriverPluginSuite covering untracked mode without a size and the normal path; the untracked test also asserts the raw key is readable so the downstream VeloxListenerApi path is covered. --- .../org/apache/gluten/GlutenPlugin.scala | 12 ++++- .../gluten/GlutenDriverPluginSuite.scala | 50 +++++++++++++++++++ 2 files changed, 61 insertions(+), 1 deletion(-) create mode 100644 gluten-core/src/test/scala/org/apache/gluten/GlutenDriverPluginSuite.scala diff --git a/gluten-core/src/main/scala/org/apache/gluten/GlutenPlugin.scala b/gluten-core/src/main/scala/org/apache/gluten/GlutenPlugin.scala index c95ddbe5b84..9cc283bd259 100644 --- a/gluten-core/src/main/scala/org/apache/gluten/GlutenPlugin.scala +++ b/gluten-core/src/main/scala/org/apache/gluten/GlutenPlugin.scala @@ -103,7 +103,8 @@ private object GlutenDriverPlugin extends Logging { } } - private def setPredefinedConfigs(conf: SparkConf): Unit = { + // Visible for testing. + private[gluten] def setPredefinedConfigs(conf: SparkConf): Unit = { // check memory off-heap enabled and size. checkOffHeapSettings(conf) @@ -134,6 +135,15 @@ private object GlutenDriverPlugin extends Logging { ((onHeapSize - (300 * 1024 * 1024)) * conf.getDouble(GlutenCoreConfig.DYNAMIC_OFFHEAP_SIZING_MEMORY_FRACTION.key, 0.6d)).toLong } else { + // Untracked memory mode skips the off-heap size requirement in checkOffHeapSettings, so + // the key may be absent here. Normalize it to 0 (mirroring the dynamic-sizing branch + // above) so downstream readers that read spark.memory.offHeap.size directly, e.g. + // VeloxListenerApi.onDriverStart, don't hit NoSuchElementException. Normal mode always has + // the key set because checkOffHeapSettings enforced it, so this only affects untracked + // mode without an explicit off-heap size. + if (!conf.contains(GlutenCoreConfig.SPARK_OFFHEAP_SIZE_KEY)) { + conf.set(GlutenCoreConfig.SPARK_OFFHEAP_SIZE_KEY, "0") + } conf.getSizeAsBytes(GlutenCoreConfig.SPARK_OFFHEAP_SIZE_KEY) } diff --git a/gluten-core/src/test/scala/org/apache/gluten/GlutenDriverPluginSuite.scala b/gluten-core/src/test/scala/org/apache/gluten/GlutenDriverPluginSuite.scala new file mode 100644 index 00000000000..9ef1f449702 --- /dev/null +++ b/gluten-core/src/test/scala/org/apache/gluten/GlutenDriverPluginSuite.scala @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.gluten + +import org.apache.gluten.config.GlutenCoreConfig + +import org.apache.spark.SparkConf + +import org.scalatest.funsuite.AnyFunSuite + +class GlutenDriverPluginSuite extends AnyFunSuite { + + test("setPredefinedConfigs does not throw in untracked mode without an off-heap size") { + // Untracked memory mode skips the off-heap size requirement in checkOffHeapSettings, so + // spark.memory.offHeap.size may be absent. setPredefinedConfigs must not read it with the + // single-arg getSizeAsBytes, which throws NoSuchElementException on a missing key. + val conf = new SparkConf(false) + .set("spark.master", "local[1]") + .set(GlutenCoreConfig.COLUMNAR_MEMORY_UNTRACKED.key, "true") + GlutenDriverPlugin.setPredefinedConfigs(conf) + assert(conf.getLong(GlutenCoreConfig.COLUMNAR_OFFHEAP_SIZE_IN_BYTES.key, -1L) == 0L) + // The raw key is normalized to 0 so downstream readers (e.g. VeloxListenerApi.onDriverStart) + // that call the single-arg getSizeAsBytes don't hit NoSuchElementException later. + assert(conf.getSizeAsBytes(GlutenCoreConfig.SPARK_OFFHEAP_SIZE_KEY) == 0L) + } + + test("setPredefinedConfigs reads the configured off-heap size in the normal path") { + val conf = new SparkConf(false) + .set("spark.master", "local[1]") + .set(GlutenCoreConfig.SPARK_OFFHEAP_ENABLED_KEY, "true") + .set(GlutenCoreConfig.SPARK_OFFHEAP_SIZE_KEY, "512m") + GlutenDriverPlugin.setPredefinedConfigs(conf) + assert( + conf.getLong(GlutenCoreConfig.COLUMNAR_OFFHEAP_SIZE_IN_BYTES.key, -1L) == 512L * 1024 * 1024) + } +}