From b4eb323e2cf690f1da829ce3578b780ae57a31b8 Mon Sep 17 00:00:00 2001 From: "shuai.xu" Date: Thu, 27 Mar 2025 10:49:28 +0800 Subject: [PATCH 01/10] [GLUTEN-9149] [core] make JniWorkspace not depend on spark --- .../backendsapi/velox/VeloxListenerApi.scala | 11 ++++++++++- .../spark/sql/expression/UDFResolver.scala | 12 ++++++++++-- .../org/apache/gluten/jni/JniLibLoader.java | 11 ----------- .../org/apache/gluten/jni/JniWorkspace.java | 18 ++---------------- 4 files changed, 22 insertions(+), 30 deletions(-) diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala index 2aa19afe7a12..9290f64c12b9 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala @@ -45,6 +45,7 @@ import org.apache.spark.util.{SparkDirectoryUtil, SparkResourceUtil} import org.apache.commons.lang3.StringUtils +import java.util.UUID import java.util.concurrent.atomic.AtomicBoolean class VeloxListenerApi extends ListenerApi with Logging { @@ -183,7 +184,14 @@ class VeloxListenerApi extends ListenerApi with Logging { UDFMappings.loadFromSparkConf(conf) // Initial library loader. - val loader = JniWorkspace.getDefault.libLoader + val loader = + JniWorkspace + .getDefault( + SparkDirectoryUtil.get + .namespace("jni") + .mkChildDirRandomly(UUID.randomUUID.toString) + .getAbsolutePath) + .libLoader // Load shared native libraries the backend libraries depend on. SharedLibraryLoader.load(conf, loader) @@ -218,6 +226,7 @@ class VeloxListenerApi extends ListenerApi with Logging { private def shutdown(): Unit = { // TODO shutdown implementation in velox to release resources + JniLibLoader.forceUnloadAll } } diff --git a/backends-velox/src/main/scala/org/apache/spark/sql/expression/UDFResolver.scala b/backends-velox/src/main/scala/org/apache/spark/sql/expression/UDFResolver.scala index 17f661289ebc..95645da4f868 100644 --- a/backends-velox/src/main/scala/org/apache/spark/sql/expression/UDFResolver.scala +++ b/backends-velox/src/main/scala/org/apache/spark/sql/expression/UDFResolver.scala @@ -32,11 +32,12 @@ import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataType, StructField, StructType} -import org.apache.spark.util.Utils +import org.apache.spark.util.{SparkDirectoryUtil, Utils} import java.io.File import java.net.URI import java.nio.file.{Files, FileVisitOption, Paths} +import java.util.UUID import scala.collection.JavaConverters.asScalaIteratorConverter import scala.collection.mutable @@ -311,7 +312,14 @@ object UDFResolver extends Logging { // Download or copy absolute paths to JniWorkspace. val uri = Utils.resolveURI(f) val name = file.getName - val jniWorkspace = new File(JniWorkspace.getDefault.getWorkDir) + val jniWorkspace = new File( + JniWorkspace + .getDefault( + SparkDirectoryUtil.get + .namespace("jni") + .mkChildDirRandomly(UUID.randomUUID.toString) + .getAbsolutePath) + .getWorkDir) if (!file.isDirectory && !f.endsWith(LIB_EXTENSION)) { val source = Utils .doFetchFile(uri.toString, Utils.createTempDir(), name, sparkConf, hadoopConf) diff --git a/gluten-core/src/main/java/org/apache/gluten/jni/JniLibLoader.java b/gluten-core/src/main/java/org/apache/gluten/jni/JniLibLoader.java index 2d07c3d79b9a..72898823f5b3 100644 --- a/gluten-core/src/main/java/org/apache/gluten/jni/JniLibLoader.java +++ b/gluten-core/src/main/java/org/apache/gluten/jni/JniLibLoader.java @@ -18,7 +18,6 @@ import org.apache.gluten.exception.GlutenException; -import org.apache.spark.util.SparkShutdownManagerUtil; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -40,22 +39,12 @@ import java.util.Set; import java.util.Vector; -import scala.runtime.BoxedUnit; - public class JniLibLoader { private static final Logger LOG = LoggerFactory.getLogger(JniLibLoader.class); private static final Set LOADED_LIBRARY_PATHS = new HashSet<>(); private static final Set REQUIRE_UNLOAD_LIBRARY_PATHS = new LinkedHashSet<>(); - static { - SparkShutdownManagerUtil.addHookForLibUnloading( - () -> { - forceUnloadAll(); - return BoxedUnit.UNIT; - }); - } - private final String workDir; private final Set loadedLibraries = new HashSet<>(); diff --git a/gluten-core/src/main/java/org/apache/gluten/jni/JniWorkspace.java b/gluten-core/src/main/java/org/apache/gluten/jni/JniWorkspace.java index a37aac03c583..348ccf9ebe90 100644 --- a/gluten-core/src/main/java/org/apache/gluten/jni/JniWorkspace.java +++ b/gluten-core/src/main/java/org/apache/gluten/jni/JniWorkspace.java @@ -20,7 +20,6 @@ import com.google.common.base.Preconditions; import org.apache.commons.io.FileUtils; -import org.apache.spark.util.SparkDirectoryUtil; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -60,19 +59,6 @@ private JniWorkspace(String rootDir) { } } - private static JniWorkspace createDefault() { - try { - final String tempRoot = - SparkDirectoryUtil.get() - .namespace("jni") - .mkChildDirRandomly(UUID.randomUUID().toString()) - .getAbsolutePath(); - return createOrGet(tempRoot); - } catch (Exception e) { - throw new GlutenException(e); - } - } - public static void enableDebug(String debugDir) { // Preserve the JNI libraries even after process exits. // This is useful for debugging native code if the debug symbols were embedded in @@ -99,10 +85,10 @@ public static void enableDebug(String debugDir) { } } - public static JniWorkspace getDefault() { + public static JniWorkspace getDefault(String rootDir) { synchronized (DEFAULT_INSTANCE_INIT_LOCK) { if (DEFAULT_INSTANCE == null) { - DEFAULT_INSTANCE = createDefault(); + DEFAULT_INSTANCE = createOrGet(rootDir); } Preconditions.checkNotNull(DEFAULT_INSTANCE); return DEFAULT_INSTANCE; From 4462d5574c97f03ef04a2bb0fab6e53c39ff4011 Mon Sep 17 00:00:00 2001 From: "shuai.xu" Date: Thu, 27 Mar 2025 18:42:24 +0800 Subject: [PATCH 02/10] unload jni libs in ch --- .../org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala | 1 + 1 file changed, 1 insertion(+) diff --git a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala index 4b661b3c2def..f2ea120f1ad6 100644 --- a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala +++ b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala @@ -131,5 +131,6 @@ class CHListenerApi extends ListenerApi with Logging { private def shutdown(): Unit = { CHBroadcastBuildSideCache.cleanAll() CHNativeExpressionEvaluator.finalizeNative() + JniLibLoader.forceUnloadAll } } From e135c36bb0076a8517da29893e1f589ac3074d81 Mon Sep 17 00:00:00 2001 From: "shuai.xu" Date: Fri, 28 Mar 2025 16:05:37 +0800 Subject: [PATCH 03/10] use add shutdown hook for ch --- .../backendsapi/clickhouse/CHListenerApi.scala | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala index f2ea120f1ad6..5f69c80b2e89 100644 --- a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala +++ b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala @@ -36,7 +36,7 @@ import org.apache.spark.rpc.{GlutenDriverEndpoint, GlutenExecutorEndpoint} import org.apache.spark.sql.execution.datasources.GlutenWriterColumnarRules import org.apache.spark.sql.execution.datasources.v1._ import org.apache.spark.sql.utils.ExpressionUtil -import org.apache.spark.util.SparkDirectoryUtil +import org.apache.spark.util.{SparkDirectoryUtil, SparkShutdownManagerUtil} import org.apache.commons.lang3.StringUtils @@ -87,6 +87,7 @@ class CHListenerApi extends ListenerApi with Logging { val executorLibPath = conf.get(GlutenConfig.GLUTEN_EXECUTOR_LIB_PATH.key, libPath) JniLibLoader.loadFromPath(executorLibPath, true) } + CHListenerApi.addShutdownHook // Add configs import org.apache.gluten.backendsapi.clickhouse.CHConfig._ conf.setCHConfig( @@ -134,3 +135,17 @@ class CHListenerApi extends ListenerApi with Logging { JniLibLoader.forceUnloadAll } } + +object CHListenerApi { + var initialized = false + + def addShutdownHook: Unit = { + if (!initialized) { + initialized = true + SparkShutdownManagerUtil.addHookForLibUnloading( + () => { + JniLibLoader.forceUnloadAll + }) + } + } +} From 4432e3d6b02f27cf50888806566b28887eaa3253 Mon Sep 17 00:00:00 2001 From: "shuai.xu" Date: Fri, 28 Mar 2025 16:06:08 +0800 Subject: [PATCH 04/10] use add shutdown hook for ch --- .../org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala | 1 - 1 file changed, 1 deletion(-) diff --git a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala index 5f69c80b2e89..120584ce2e5e 100644 --- a/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala +++ b/backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHListenerApi.scala @@ -132,7 +132,6 @@ class CHListenerApi extends ListenerApi with Logging { private def shutdown(): Unit = { CHBroadcastBuildSideCache.cleanAll() CHNativeExpressionEvaluator.finalizeNative() - JniLibLoader.forceUnloadAll } } From be45858b25abb60eef93b3503cd18dafffa579d0 Mon Sep 17 00:00:00 2001 From: "shuai.xu" Date: Mon, 31 Mar 2025 17:33:46 +0800 Subject: [PATCH 05/10] address comments --- .../backendsapi/velox/VeloxListenerApi.scala | 18 ++++++++---------- .../spark/sql/expression/UDFResolver.scala | 12 ++---------- .../org/apache/gluten/jni/JniWorkspace.java | 9 +++++++-- 3 files changed, 17 insertions(+), 22 deletions(-) diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala index 9290f64c12b9..4ef9aa8c7f92 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala @@ -117,7 +117,6 @@ class VeloxListenerApi extends ListenerApi with Logging { } SparkDirectoryUtil.init(conf) - UDFResolver.resolveUdfConf(conf, isDriver = true) initialize(conf, isDriver = true) UdfJniWrapper.registerFunctionSignatures() } @@ -144,13 +143,19 @@ class VeloxListenerApi extends ListenerApi with Logging { } SparkDirectoryUtil.init(conf) - UDFResolver.resolveUdfConf(conf, isDriver = false) initialize(conf, isDriver = false) } override def onExecutorShutdown(): Unit = shutdown() private def initialize(conf: SparkConf, isDriver: Boolean): Unit = { + JniWorkspace.intializeDefault( + SparkDirectoryUtil.get + .namespace("jni") + .mkChildDirRandomly(UUID.randomUUID.toString) + .getAbsolutePath) + UDFResolver.resolveUdfConf(conf, isDriver) + // Do row / batch type initializations. Convention.ensureSparkRowAndBatchTypesRegistered() ArrowJavaBatch.ensureRegistered() @@ -184,14 +189,7 @@ class VeloxListenerApi extends ListenerApi with Logging { UDFMappings.loadFromSparkConf(conf) // Initial library loader. - val loader = - JniWorkspace - .getDefault( - SparkDirectoryUtil.get - .namespace("jni") - .mkChildDirRandomly(UUID.randomUUID.toString) - .getAbsolutePath) - .libLoader + val loader = JniWorkspace.getDefault.libLoader // Load shared native libraries the backend libraries depend on. SharedLibraryLoader.load(conf, loader) diff --git a/backends-velox/src/main/scala/org/apache/spark/sql/expression/UDFResolver.scala b/backends-velox/src/main/scala/org/apache/spark/sql/expression/UDFResolver.scala index 95645da4f868..17f661289ebc 100644 --- a/backends-velox/src/main/scala/org/apache/spark/sql/expression/UDFResolver.scala +++ b/backends-velox/src/main/scala/org/apache/spark/sql/expression/UDFResolver.scala @@ -32,12 +32,11 @@ import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DataType, StructField, StructType} -import org.apache.spark.util.{SparkDirectoryUtil, Utils} +import org.apache.spark.util.Utils import java.io.File import java.net.URI import java.nio.file.{Files, FileVisitOption, Paths} -import java.util.UUID import scala.collection.JavaConverters.asScalaIteratorConverter import scala.collection.mutable @@ -312,14 +311,7 @@ object UDFResolver extends Logging { // Download or copy absolute paths to JniWorkspace. val uri = Utils.resolveURI(f) val name = file.getName - val jniWorkspace = new File( - JniWorkspace - .getDefault( - SparkDirectoryUtil.get - .namespace("jni") - .mkChildDirRandomly(UUID.randomUUID.toString) - .getAbsolutePath) - .getWorkDir) + val jniWorkspace = new File(JniWorkspace.getDefault.getWorkDir) if (!file.isDirectory && !f.endsWith(LIB_EXTENSION)) { val source = Utils .doFetchFile(uri.toString, Utils.createTempDir(), name, sparkConf, hadoopConf) diff --git a/gluten-core/src/main/java/org/apache/gluten/jni/JniWorkspace.java b/gluten-core/src/main/java/org/apache/gluten/jni/JniWorkspace.java index 348ccf9ebe90..a9677352ba2c 100644 --- a/gluten-core/src/main/java/org/apache/gluten/jni/JniWorkspace.java +++ b/gluten-core/src/main/java/org/apache/gluten/jni/JniWorkspace.java @@ -85,12 +85,17 @@ public static void enableDebug(String debugDir) { } } - public static JniWorkspace getDefault(String rootDir) { + public static void initializeDefault(String rootDir) { synchronized (DEFAULT_INSTANCE_INIT_LOCK) { if (DEFAULT_INSTANCE == null) { DEFAULT_INSTANCE = createOrGet(rootDir); } - Preconditions.checkNotNull(DEFAULT_INSTANCE); + } + } + + public static JniWorkspace getDefault() { + synchronized (DEFAULT_INSTANCE_INIT_LOCK) { + Preconditions.checkNotNull(DEFAULT_INSTANCE, "Not call initializeDefault yet"); return DEFAULT_INSTANCE; } } From d5d62eceb7eb66bbece1cd192b3dee179e4623aa Mon Sep 17 00:00:00 2001 From: "shuai.xu" Date: Mon, 31 Mar 2025 17:46:17 +0800 Subject: [PATCH 06/10] fix build break --- .../org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala index 4ef9aa8c7f92..2c8d0bc668a1 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala @@ -149,7 +149,7 @@ class VeloxListenerApi extends ListenerApi with Logging { override def onExecutorShutdown(): Unit = shutdown() private def initialize(conf: SparkConf, isDriver: Boolean): Unit = { - JniWorkspace.intializeDefault( + JniWorkspace.initializeDefault( SparkDirectoryUtil.get .namespace("jni") .mkChildDirRandomly(UUID.randomUUID.toString) From 80771b28d88beeb8323673a298dbf2c8c6dc8b08 Mon Sep 17 00:00:00 2001 From: "shuai.xu" Date: Thu, 3 Apr 2025 11:02:40 +0800 Subject: [PATCH 07/10] address comments and fix a bug --- .../backendsapi/velox/VeloxListenerApi.scala | 22 ++++++++++--------- .../org/apache/gluten/jni/JniWorkspace.java | 11 +++++----- 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala index 2c8d0bc668a1..0c6c0fdf5b64 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala @@ -149,11 +149,19 @@ class VeloxListenerApi extends ListenerApi with Logging { override def onExecutorShutdown(): Unit = shutdown() private def initialize(conf: SparkConf, isDriver: Boolean): Unit = { + // Sets this configuration only once, since not undoable. + // DebugInstance should be created first. + if (conf.getBoolean(GlutenConfig.DEBUG_KEEP_JNI_WORKSPACE.key, defaultValue = false)) { + val debugDir = conf.get(GlutenConfig.DEBUG_KEEP_JNI_WORKSPACE_DIR.key) + JniWorkspace.enableDebug(debugDir) + } JniWorkspace.initializeDefault( - SparkDirectoryUtil.get - .namespace("jni") - .mkChildDirRandomly(UUID.randomUUID.toString) - .getAbsolutePath) + () => + SparkDirectoryUtil.get + .namespace("jni") + .mkChildDirRandomly(UUID.randomUUID.toString) + .getAbsolutePath) + UDFResolver.resolveUdfConf(conf, isDriver) // Do row / batch type initializations. @@ -175,12 +183,6 @@ class VeloxListenerApi extends ListenerApi with Logging { classOf[ColumnarShuffleManager].getName ) - // Sets this configuration only once, since not undoable. - if (conf.getBoolean(GlutenConfig.DEBUG_KEEP_JNI_WORKSPACE.key, defaultValue = false)) { - val debugDir = conf.get(GlutenConfig.DEBUG_KEEP_JNI_WORKSPACE_DIR.key) - JniWorkspace.enableDebug(debugDir) - } - // Set the system properties. // Use appending policy for children with the same name in a arrow struct vector. System.setProperty("arrow.struct.conflict.policy", "CONFLICT_APPEND") diff --git a/gluten-core/src/main/java/org/apache/gluten/jni/JniWorkspace.java b/gluten-core/src/main/java/org/apache/gluten/jni/JniWorkspace.java index a9677352ba2c..50827c7a0185 100644 --- a/gluten-core/src/main/java/org/apache/gluten/jni/JniWorkspace.java +++ b/gluten-core/src/main/java/org/apache/gluten/jni/JniWorkspace.java @@ -31,6 +31,7 @@ import java.util.Map; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Supplier; public class JniWorkspace { private static final Logger LOG = LoggerFactory.getLogger(JniWorkspace.class); @@ -85,19 +86,17 @@ public static void enableDebug(String debugDir) { } } - public static void initializeDefault(String rootDir) { + public static void initializeDefault(Supplier rootDir) { synchronized (DEFAULT_INSTANCE_INIT_LOCK) { if (DEFAULT_INSTANCE == null) { - DEFAULT_INSTANCE = createOrGet(rootDir); + DEFAULT_INSTANCE = createOrGet(rootDir.get()); } } } public static JniWorkspace getDefault() { - synchronized (DEFAULT_INSTANCE_INIT_LOCK) { - Preconditions.checkNotNull(DEFAULT_INSTANCE, "Not call initializeDefault yet"); - return DEFAULT_INSTANCE; - } + Preconditions.checkNotNull(DEFAULT_INSTANCE, "Not call initializeDefault yet"); + return DEFAULT_INSTANCE; } private static JniWorkspace createOrGet(String rootDir) { From 8362cffd9a5c8738037c7dda61dabff223b79e91 Mon Sep 17 00:00:00 2001 From: "shuai.xu" Date: Thu, 3 Apr 2025 15:43:24 +0800 Subject: [PATCH 08/10] trigger ci --- .../gluten/backendsapi/velox/VeloxListenerApi.scala | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala index 0c6c0fdf5b64..bc0b62ca7c8d 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala @@ -154,13 +154,14 @@ class VeloxListenerApi extends ListenerApi with Logging { if (conf.getBoolean(GlutenConfig.DEBUG_KEEP_JNI_WORKSPACE.key, defaultValue = false)) { val debugDir = conf.get(GlutenConfig.DEBUG_KEEP_JNI_WORKSPACE_DIR.key) JniWorkspace.enableDebug(debugDir) + } else { + JniWorkspace.initializeDefault( + () => + SparkDirectoryUtil.get + .namespace("jni") + .mkChildDirRandomly(UUID.randomUUID.toString) + .getAbsolutePath) } - JniWorkspace.initializeDefault( - () => - SparkDirectoryUtil.get - .namespace("jni") - .mkChildDirRandomly(UUID.randomUUID.toString) - .getAbsolutePath) UDFResolver.resolveUdfConf(conf, isDriver) From 33fdd376538d408920d146073b9ec095135df521 Mon Sep 17 00:00:00 2001 From: "shuai.xu" Date: Mon, 7 Apr 2025 09:32:03 +0800 Subject: [PATCH 09/10] remove unload lib --- .../org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala | 1 - 1 file changed, 1 deletion(-) diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala index bc0b62ca7c8d..8167a3b2dadb 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala @@ -227,7 +227,6 @@ class VeloxListenerApi extends ListenerApi with Logging { private def shutdown(): Unit = { // TODO shutdown implementation in velox to release resources - JniLibLoader.forceUnloadAll } } From 1d22fa4544301b9d2ca29ed7c804c30302b7ffb5 Mon Sep 17 00:00:00 2001 From: "shuai.xu" Date: Mon, 7 Apr 2025 12:27:40 +0800 Subject: [PATCH 10/10] add back shutdown hook --- .../gluten/backendsapi/velox/VeloxListenerApi.scala | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala index 8167a3b2dadb..109f3015a691 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxListenerApi.scala @@ -41,7 +41,7 @@ import org.apache.spark.sql.execution.datasources.GlutenWriterColumnarRules import org.apache.spark.sql.execution.datasources.velox.{VeloxParquetWriterInjects, VeloxRowSplitter} import org.apache.spark.sql.expression.UDFResolver import org.apache.spark.sql.internal.{GlutenConfigUtil, StaticSQLConf} -import org.apache.spark.util.{SparkDirectoryUtil, SparkResourceUtil} +import org.apache.spark.util.{SparkDirectoryUtil, SparkResourceUtil, SparkShutdownManagerUtil} import org.apache.commons.lang3.StringUtils @@ -149,6 +149,7 @@ class VeloxListenerApi extends ListenerApi with Logging { override def onExecutorShutdown(): Unit = shutdown() private def initialize(conf: SparkConf, isDriver: Boolean): Unit = { + addShutdownHook // Sets this configuration only once, since not undoable. // DebugInstance should be created first. if (conf.getBoolean(GlutenConfig.DEBUG_KEEP_JNI_WORKSPACE.key, defaultValue = false)) { @@ -250,4 +251,11 @@ object VeloxListenerApi { private def inLocalMode(conf: SparkConf): Boolean = { SparkResourceUtil.isLocalMaster(conf) } + + private def addShutdownHook: Unit = { + SparkShutdownManagerUtil.addHookForLibUnloading( + () => { + JniLibLoader.forceUnloadAll + }) + } }