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

common/utils/src/main/scala/org/apache/spark/util/SparkEnvUtils.scala

@@ -27,6 +27,10 @@ private[spark] trait SparkEnvUtils {
     System.getenv("SPARK_TESTING") != null || System.getProperty("spark.testing") != null
   }
 
+  /**
+   * Whether allow using native BLAS/LAPACK/ARPACK libraries if available.
+   */
+  val allowNativeBlas = "true".equals(System.getProperty("netlib.allowNativeBlas", "true"))
 }
 
 object SparkEnvUtils extends SparkEnvUtils

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)
 
@@ -3409,32 +3410,35 @@ object SparkContext extends Logging {
     }
   }
 
+  private def supplementJavaOpts(
+      conf: SparkConf, key: OptionalConfigEntry[String], javaOpts: String): Unit = {
+    val v = conf.get(key) match {
+      case Some(opts) => s"$javaOpts $opts"
+      case None => javaOpts
+    }
+    conf.set(key.key, v)
+  }
+
   /**
    * SPARK-36796: This is a helper function to supplement some JVM runtime options to
    * `spark.driver.extraJavaOptions` and `spark.executor.extraJavaOptions`.
    */
   private def supplementJavaModuleOptions(conf: SparkConf): Unit = {
-    def supplement(key: OptionalConfigEntry[String]): Unit = {
-      val v = conf.get(key) match {
-        case Some(opts) => s"${JavaModuleOptions.defaultModuleOptions()} $opts"
-        case None => JavaModuleOptions.defaultModuleOptions()
-      }
-      conf.set(key.key, v)
-    }
-    supplement(DRIVER_JAVA_OPTIONS)
-    supplement(EXECUTOR_JAVA_OPTIONS)
+    val opts = JavaModuleOptions.defaultModuleOptions()
+    supplementJavaOpts(conf, DRIVER_JAVA_OPTIONS, opts)
+    supplementJavaOpts(conf, EXECUTOR_JAVA_OPTIONS, opts)
   }
 
   private def supplementJavaIPv6Options(conf: SparkConf): Unit = {
-    def supplement(key: OptionalConfigEntry[String]): Unit = {
-      val v = conf.get(key) match {
-        case Some(opts) => s"-Djava.net.preferIPv6Addresses=${Utils.preferIPv6} $opts"
-        case None => s"-Djava.net.preferIPv6Addresses=${Utils.preferIPv6}"
-      }
-      conf.set(key.key, v)
-    }
-    supplement(DRIVER_JAVA_OPTIONS)
-    supplement(EXECUTOR_JAVA_OPTIONS)
+    val opts = s"-Djava.net.preferIPv6Addresses=${Utils.preferIPv6}"
+    supplementJavaOpts(conf, DRIVER_JAVA_OPTIONS, opts)
+    supplementJavaOpts(conf, EXECUTOR_JAVA_OPTIONS, opts)
+  }
+
+  private def supplementBlasOptions(conf: SparkConf): Unit = {
+    val opts = s"-Dnetlib.allowNativeBlas=${Utils.allowNativeBlas}"
+    supplementJavaOpts(conf, DRIVER_JAVA_OPTIONS, opts)
+    supplementJavaOpts(conf, EXECUTOR_JAVA_OPTIONS, opts)
   }
 }
 

core/src/main/scala/org/apache/spark/internal/config/package.scala

@@ -2833,4 +2833,12 @@ package object config {
       .checkValues(Set("connect", "classic"))
       .createWithDefault(
         if (sys.env.get("SPARK_CONNECT_MODE").contains("1")) "connect" else "classic")
+
+  private[spark] val SPARK_ML_ALLOW_NATIVE_BLAS =
+    ConfigBuilder("spark.ml.allowNativeBlas")
+      .doc("Whether allow using native BLAS/LAPACK/ARPACK implementations when native " +
+        "libraries are available. If disabled, always use Java implementations.")
+      .version("4.1.0")
+      .booleanConf
+      .createWithDefault(true)
 }

docs/ml-linalg-guide.md

@@ -46,8 +46,7 @@ The installation should be done on all nodes of the cluster. Generic version of
 
 For Debian / Ubuntu:
 ```
-sudo apt-get install libopenblas-base
-sudo update-alternatives --config libblas.so.3
+sudo apt-get install libopenblas-dev
 ```
 For CentOS / RHEL:
 ```
@@ -76,6 +75,8 @@ You can also point `dev.ludovic.netlib` to specific libraries names and paths. F
 
 If native libraries are not properly configured in the system, the Java implementation (javaBLAS) will be used as fallback option.
 
+You can also set spark conf `spark.ml.allowNativeBlas` or Java system property `netlib.allowNativeBlas` to `false` to disable native BLAS and always use the Java implementation.
+
 ## Spark Configuration
 
 The default behavior of multi-threading in either Intel MKL or OpenBLAS may not be optimal with Spark's execution model [^1].

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 allowNativeBlas = config.get("spark.ml.allowNativeBlas");
+      addOptionString(cmd, "-Dnetlib.allowNativeBlas=" + allowNativeBlas);
+    }
+
     // SPARK-36796: Always add some JVM runtime default options to submit command
     addOptionString(cmd, JavaModuleOptions.defaultModuleOptions());
     addOptionString(cmd, "-Dderby.connection.requireAuthentication=false");

mllib-local/pom.xml

@@ -66,6 +66,11 @@
       <groupId>org.apache.spark</groupId>
       <artifactId>spark-tags_${scala.binary.version}</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-common-utils_${scala.binary.version}</artifactId>
+      <version>${project.version}</version>
+    </dependency>
 
     <!--
       This spark-tags test-dep is needed even though it isn't used in this module, otherwise testing-cmds that exclude

mllib-local/src/main/scala/org/apache/spark/ml/linalg/BLAS.scala

@@ -19,10 +19,13 @@ package org.apache.spark.ml.linalg
 
 import dev.ludovic.netlib.blas.{BLAS => NetlibBLAS, JavaBLAS => NetlibJavaBLAS, NativeBLAS => NetlibNativeBLAS}
 
+import org.apache.spark.internal.Logging
+import org.apache.spark.util.SparkEnvUtils
+
 /**
  * BLAS routines for MLlib's vectors and matrices.
  */
-private[spark] object BLAS extends Serializable {
+private[spark] object BLAS extends Serializable with Logging {
 
   @transient private var _javaBLAS: NetlibBLAS = _
   @transient private var _nativeBLAS: NetlibBLAS = _
@@ -39,8 +42,12 @@ private[spark] object BLAS extends Serializable {
   // For level-3 routines, we use the native BLAS.
   private[spark] def nativeBLAS: NetlibBLAS = {
     if (_nativeBLAS == null) {
-      _nativeBLAS =
+      _nativeBLAS = if (SparkEnvUtils.allowNativeBlas) {
         try { NetlibNativeBLAS.getInstance } catch { case _: Throwable => javaBLAS }
+      } else {
+        logInfo("Disable native BLAS because netlib.allowNativeBlas is false.")
+        javaBLAS
+      }
     }
     _nativeBLAS
   }

mllib/src/main/scala/org/apache/spark/mllib/linalg/ARPACK.scala

@@ -19,10 +19,13 @@ package org.apache.spark.mllib.linalg
 
 import dev.ludovic.netlib.arpack.{ARPACK => NetlibARPACK, JavaARPACK => NetlibJavaARPACK, NativeARPACK => NetlibNativeARPACK}
 
+import org.apache.spark.internal.Logging
+import org.apache.spark.util.SparkEnvUtils
+
 /**
  * ARPACK routines for MLlib's vectors and matrices.
  */
-private[spark] object ARPACK extends Serializable {
+private[spark] object ARPACK extends Serializable with Logging {
 
   @transient private var _javaARPACK: NetlibARPACK = _
   @transient private var _nativeARPACK: NetlibARPACK = _
@@ -36,8 +39,12 @@ private[spark] object ARPACK extends Serializable {
 
   private[spark] def nativeARPACK: NetlibARPACK = {
     if (_nativeARPACK == null) {
-      _nativeARPACK =
+      _nativeARPACK = if (SparkEnvUtils.allowNativeBlas) {
         try { NetlibNativeARPACK.getInstance } catch { case _: Throwable => javaARPACK }
+      } else {
+        logInfo("Disable native ARPACK because netlib.allowNativeBlas is false.")
+        javaARPACK
+      }
     }
     _nativeARPACK
   }

mllib/src/main/scala/org/apache/spark/mllib/linalg/BLAS.scala

@@ -17,44 +17,14 @@
 
 package org.apache.spark.mllib.linalg
 
-import dev.ludovic.netlib.blas.{BLAS => NetlibBLAS, JavaBLAS => NetlibJavaBLAS, NativeBLAS => NetlibNativeBLAS}
-
 import org.apache.spark.internal.Logging
+import org.apache.spark.ml.linalg.BLAS.{getBLAS, nativeBLAS}
 
 /**
  * BLAS routines for MLlib's vectors and matrices.
  */
 private[spark] object BLAS extends Serializable with Logging {
 
-  @transient private var _javaBLAS: NetlibBLAS = _
-  @transient private var _nativeBLAS: NetlibBLAS = _
-  private val nativeL1Threshold: Int = 256
-
-  // For level-1 function dspmv, use javaBLAS for better performance.
-  private[spark] def javaBLAS: NetlibBLAS = {
-    if (_javaBLAS == null) {
-      _javaBLAS = NetlibJavaBLAS.getInstance
-    }
-    _javaBLAS
-  }
-
-  // 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
-  }
-
-  private[spark] def getBLAS(vectorSize: Int): NetlibBLAS = {
-    if (vectorSize < nativeL1Threshold) {
-      javaBLAS
-    } else {
-      nativeBLAS
-    }
-  }
-
   /**
    * y += a * x
    */

mllib/src/main/scala/org/apache/spark/mllib/linalg/LAPACK.scala

@@ -19,10 +19,13 @@ package org.apache.spark.mllib.linalg
 
 import dev.ludovic.netlib.lapack.{JavaLAPACK => NetlibJavaLAPACK, LAPACK => NetlibLAPACK, NativeLAPACK => NetlibNativeLAPACK}
 
+import org.apache.spark.internal.Logging
+import org.apache.spark.util.SparkEnvUtils
+
 /**
  * LAPACK routines for MLlib's vectors and matrices.
  */
-private[spark] object LAPACK extends Serializable {
+private[spark] object LAPACK extends Serializable with Logging {
 
   @transient private var _javaLAPACK: NetlibLAPACK = _
   @transient private var _nativeLAPACK: NetlibLAPACK = _
@@ -36,8 +39,12 @@ private[spark] object LAPACK extends Serializable {
 
   private[spark] def nativeLAPACK: NetlibLAPACK = {
     if (_nativeLAPACK == null) {
-      _nativeLAPACK =
+      _nativeLAPACK = if (SparkEnvUtils.allowNativeBlas) {
         try { NetlibNativeLAPACK.getInstance } catch { case _: Throwable => javaLAPACK }
+      } else {
+        logInfo("Disable native LAPACK because netlib.allowNativeBlas is false.")
+        javaLAPACK
+      }
     }
     _nativeLAPACK
   }

mllib/src/main/scala/org/apache/spark/mllib/linalg/Matrices.scala

@@ -427,7 +427,7 @@ class DenseMatrix @Since("1.3.0") (
     if (isTransposed) {
       Iterator.tabulate(numCols) { j =>
         val col = new Array[Double](numRows)
-        BLAS.nativeBLAS.dcopy(numRows, values, j, numCols, col, 0, 1)
+        newlinalg.BLAS.nativeBLAS.dcopy(numRows, values, j, numCols, col, 0, 1)
         new DenseVector(col)
       }
     } else {

mllib/src/main/scala/org/apache/spark/mllib/recommendation/MatrixFactorizationModel.scala

@@ -31,7 +31,7 @@ import org.apache.spark.SparkContext
 import org.apache.spark.annotation.Since
 import org.apache.spark.api.java.{JavaPairRDD, JavaRDD}
 import org.apache.spark.internal.{Logging, LogKeys, MDC}
-import org.apache.spark.mllib.linalg.BLAS
+import org.apache.spark.ml.linalg.BLAS
 import org.apache.spark.mllib.rdd.MLPairRDDFunctions._
 import org.apache.spark.mllib.util.{Loader, Saveable}
 import org.apache.spark.rdd.RDD

mllib/src/test/scala/org/apache/spark/mllib/linalg/BLASSuite.scala

@@ -23,12 +23,6 @@ import org.apache.spark.mllib.util.TestingUtils._
 
 class BLASSuite extends SparkFunSuite {
 
-  test("nativeL1Threshold") {
-    assert(getBLAS(128) == BLAS.javaBLAS)
-    assert(getBLAS(256) == BLAS.nativeBLAS)
-    assert(getBLAS(512) == BLAS.nativeBLAS)
-  }
-
   test("copy") {
     val sx = Vectors.sparse(4, Array(0, 2), Array(1.0, -2.0))
     val dx = Vectors.dense(1.0, 0.0, -2.0, 0.0)

resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala

@@ -1049,6 +1049,8 @@ private[spark] class Client(
 
     javaOpts += s"-Djava.net.preferIPv6Addresses=${Utils.preferIPv6}"
 
+    javaOpts += s"-Dnetlib.allowNativeBlas=${sparkConf.get(SPARK_ML_ALLOW_NATIVE_BLAS)}"
+
     // SPARK-37106: To start AM with Java 17, `JavaModuleOptions.defaultModuleOptions`
     // is added by default.
     javaOpts += JavaModuleOptions.defaultModuleOptions()