[SPARK-51243][CORE][ML] Configurable allow native BLAS

apache · Feb 17, 2025 · 265f798 · 265f798
1 parent e1f7851
commit 265f798
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 2 deletions.
diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -432,6 +432,7 @@ class SparkContext(config: SparkConf) extends Logging {
 
     SparkContext.supplementJavaModuleOptions(_conf)
     SparkContext.supplementJavaIPv6Options(_conf)
+    SparkContext.supplementBlasOptions(_conf)
 
     _driverLogger = DriverLogger(_conf)
 
@@ -3436,6 +3437,20 @@ object SparkContext extends Logging {
     supplement(DRIVER_JAVA_OPTIONS)
     supplement(EXECUTOR_JAVA_OPTIONS)
   }
+
+  private def supplementCustomOptions(conf: SparkConf): Unit = {
+    conf.getOption("spark.ml.allowNativeBlas").foreach { allowNativeBlas =>
+      def supplement(key: OptionalConfigEntry[String]): Unit = {
+        val v = conf.get(key) match {
+          case Some(opts) => s"-Dspark.ml.allowNativeBlas=$allowNativeBlas $opts"
+          case None => s"-Dspark.ml.allowNativeBlas=$allowNativeBlas"
+        }
+        conf.set(key.key, v)
+      }
+      supplement(DRIVER_JAVA_OPTIONS)
+      supplement(EXECUTOR_JAVA_OPTIONS)
+    }
+  }
 }
 
 /**

diff --git a/launcher/src/main/java/org/apache/spark/launcher/SparkSubmitCommandBuilder.java b/launcher/src/main/java/org/apache/spark/launcher/SparkSubmitCommandBuilder.java
@@ -332,6 +332,11 @@ private List<String> buildSparkSubmitCommand(Map<String, String> env)
         config.get(SparkLauncher.DRIVER_EXTRA_LIBRARY_PATH));
     }
 
+    if (config.containsKey("spark.ml.allowNativeBlas")) {
+      String preferNativeBlas = config.get("spark.ml.allowNativeBlas");
+      addOptionString(cmd, "-Dspark.ml.allowNativeBlas=" + preferNativeBlas);
+    }
+
     // SPARK-36796: Always add some JVM runtime default options to submit command
     addOptionString(cmd, JavaModuleOptions.defaultModuleOptions());
     addOptionString(cmd, "-Dderby.connection.requireAuthentication=false");

diff --git a/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala b/mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala
@@ -39,8 +39,11 @@ private[spark] object BLAS extends Serializable {
   // For level-3 routines, we use the native BLAS.
   private[spark] def nativeBLAS: NetlibBLAS = {
     if (_nativeBLAS == null) {
-      _nativeBLAS =
-        try { NetlibNativeBLAS.getInstance } catch { case _: Throwable => javaBLAS }
+      _nativeBLAS = System.getProperty("spark.ml.allowNativeBlas", "true") match {
+        case "true" =>
+          try { NetlibNativeBLAS.getInstance } catch { case _: Throwable => javaBLAS }
+        case _ => javaBLAS
+      }
     }
     _nativeBLAS
   }

diff --git a/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala b/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
@@ -1049,6 +1049,10 @@ private[spark] class Client(
 
     javaOpts += s"-Djava.net.preferIPv6Addresses=${Utils.preferIPv6}"
 
+    sparkConf.getOption("spark.ml.allowNativeBlas").foreach { allowNativeBlas =>
+      javaOpts += s"-Dspark.ml.allowNativeBlas=$allowNativeBlas"
+    }
+
     // SPARK-37106: To start AM with Java 17, `JavaModuleOptions.defaultModuleOptions`
     // is added by default.
     javaOpts += JavaModuleOptions.defaultModuleOptions()