You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by mg...@apache.org on 2025/01/16 14:16:06 UTC
(spark) branch master updated: [SPARK-50827][ML][CONNECT] Support plugin
This is an automated email from the ASF dual-hosted git repository.
mgrund pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 86804f6da5fc [SPARK-50827][ML][CONNECT] Support plugin
86804f6da5fc is described below
commit 86804f6da5fc0ad1b1a79864d99c31a5afff8f7a
Author: Bobby Wang <wb...@gmail.com>
AuthorDate: Thu Jan 16 15:15:50 2025 +0100
[SPARK-50827][ML][CONNECT] Support plugin
### What changes were proposed in this pull request?
This PR introduces connect ML with a plugin that could replace operators of spark.ml with third-party implementations.
### Why are the changes needed?
The connect ML is using spark.ml as the backend implementations, which is not GPU-accelerated. This PR could allow the
third-party library as the connect ML backend to accelerate the ML cases with GPUs without changing any users code.
### Does this PR introduce _any_ user-facing change?
Yes, it's a new feature. This PR defines a spark configuration `spark.connect.ml.backend.classes` to specify the plugin which
must implement `org.apache.spark.sql.connect.plugin.MLBackendPlugin`.
### How was this patch tested?
Make sure the CI (the newly added tests which have covered the plugin cases) pass. And manually run below code without any exception.
1. Build your own plugin package
- Customize the estimator
``` scala
package com.example.ml
import org.apache.commons.logging.LogFactory
import org.apache.spark.ml.classification.{LogisticRegression => SparkLR, LogisticRegressionModel => SparkLRModel}
import org.apache.spark.ml.util.Identifiable
import org.apache.spark.sql.Dataset
class LogisticRegression(override val uid: String) extends SparkLR {
private val logger = LogFactory.getLog("com.example.ml.LogisticRegression")
def this() = this(Identifiable.randomUID("logreg"))
override def train(dataset: Dataset[_]): SparkLRModel = {
logger.info("Bobby train in SparkRapidsML library.")
super.train(dataset)
}
}
```
and create `org.apache.spark.ml.Estimator` file under the `META-INF/services` and register the customized LogisticRegression
```
com.example.ml.LogisticRegression
```
- Implement the MLBackendPlugin
``` scala
package com.example.ml
import org.apache.spark.sql.connect.plugin.MLBackendPlugin
import java.util.Optional
class Plugin extends MLBackendPlugin {
override def transform(mlName: String): Optional[String] = {
mlName match {
case "org.apache.spark.ml.classification.LogisticRegression" =>
Optional.of("com.example.ml.LogisticRegression")
case _ => Optional.empty()
}
}
}
```
- Compile
After compilation, we can get a similiar com.example.ml-1.0-SNAPSHOT.jar
2. Run the python tests
``` python
from pyspark.ml.classification import (LogisticRegression,
LogisticRegressionModel)
from pyspark.ml.linalg import Vectors
from pyspark.sql import SparkSession
os.environ["PYSPARK_PYTHON"] = "/home/xxx/anaconda3/envs/pyspark/bin/python"
os.environ["PYSPARK_DRIVER_PYTHON"] = "/home/xxx/anaconda3/envs/pyspark/bin/python"
os.environ["SPARK_CONNECT_MODE_ENABLED"] = ""
spark = (SparkSession.builder.remote("sc://localhost")
.config("spark.connect.ml.backend.classes", "com.example.ml.Plugin")
.getOrCreate())
spark.addArtifact("target/com.example.ml-1.0-SNAPSHOT.jar")
def run_test():
df = spark.createDataFrame([
(Vectors.dense([1.0, 2.0]), 1),
(Vectors.dense([2.0, -1.0]), 1),
(Vectors.dense([-3.0, -2.0]), 0),
(Vectors.dense([-1.0, -2.0]), 0),
], schema=['features', 'label'])
lr = LogisticRegression()
lr.setMaxIter(30)
lr.setThreshold(0.8)
model: LogisticRegressionModel = lr.fit(df)
assert model.getMaxIter() == 30
assert model.getThreshold() == 0.8
print(f"model: {model}")
x = model.predictRaw(Vectors.dense([1.0, 2.0]))
print(f"predictRaw {x}")
print(f"model.coefficients: {model.coefficients}")
print(f"model.intercept: {model.intercept}")
print("-------------------- done ---------------- ")
# Train with com.example.ml.LogisticRegression
run_test()
```
Then we can check if `SPARK_HOME/logs/**org.apache.spark.sql.connect.service.SparkConnectServer-1**` contains `25/01/15 15:49:27 INFO LogisticRegression: Bobby train in SparkRapidsML library.`.
### Was this patch authored or co-authored using generative AI tooling?
No
Closes #49503 from wbo4958/plugin.connect.ml.
Authored-by: Bobby Wang <wb...@gmail.com>
Signed-off-by: Martin Grund <ma...@databricks.com>
---
.../spark/sql/connect/plugin/MLBackendPlugin.java | 37 ++++
.../apache/spark/sql/connect/config/Connect.scala | 12 ++
.../apache/spark/sql/connect/ml/MLHandler.scala | 14 +-
.../org/apache/spark/sql/connect/ml/MLUtils.scala | 62 +++++--
.../plugin/SparkConnectPluginRegistry.scala | 6 +
.../services/org.apache.spark.ml.Estimator | 20 +++
.../spark/sql/connect/ml/MLBackendSuite.scala | 193 +++++++++++++++++++++
.../org/apache/spark/sql/connect/ml/MLHelper.scala | 138 +++++++++++++++
.../org/apache/spark/sql/connect/ml/MLSuite.scala | 156 ++++++++---------
9 files changed, 540 insertions(+), 98 deletions(-)
diff --git a/sql/connect/server/src/main/java/org/apache/spark/sql/connect/plugin/MLBackendPlugin.java b/sql/connect/server/src/main/java/org/apache/spark/sql/connect/plugin/MLBackendPlugin.java
new file mode 100644
index 000000000000..a7b7d267c0fb
--- /dev/null
+++ b/sql/connect/server/src/main/java/org/apache/spark/sql/connect/plugin/MLBackendPlugin.java
@@ -0,0 +1,37 @@
+/*
+ * 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.spark.sql.connect.plugin;
+
+import java.util.Optional;
+
+/**
+ * :: Experimental ::
+ * Behavior trait for supporting backend for the Spark Connect ML.
+ *
+ * @since 4.0.0
+ */
+public interface MLBackendPlugin {
+ /**
+ * Replaces the specified Spark ML operators with a backend-specific implementation.
+ * Returns Optional.empty() if no replacement is provided for the given operator
+ *
+ * @param mlName The name of the operator
+ * @return an optional replaced operator name
+ */
+ Optional<String> transform(String mlName);
+}
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala
index b64637f7d247..b0c5a2a055b5 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala
@@ -214,6 +214,18 @@ object Connect {
.toSequence
.createWithDefault(Nil)
+ val CONNECT_ML_BACKEND_CLASSES =
+ buildConf("spark.connect.ml.backend.classes")
+ .doc("""
+ |Comma separated list of classes that implement the trait
+ |org.apache.spark.sql.connect.plugin.MLBackendPlugin to replace the
+ |specified Spark ML operators with a backend-specific implementation.
+ |""".stripMargin)
+ .version("4.0.0")
+ .stringConf
+ .toSequence
+ .createWithDefault(Nil)
+
val CONNECT_JVM_STACK_TRACE_MAX_SIZE =
buildStaticConf("spark.connect.jvmStacktrace.maxSize")
.doc("""
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala
index b4bc6bfdc66b..1f0981e048af 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala
@@ -117,7 +117,8 @@ private[connect] object MLHandler extends Logging {
assert(estimatorProto.getType == proto.MlOperator.OperatorType.ESTIMATOR)
val dataset = MLUtils.parseRelationProto(fitCmd.getDataset, sessionHolder)
- val estimator = MLUtils.getEstimator(estimatorProto, Some(fitCmd.getParams))
+ val estimator =
+ MLUtils.getEstimator(sessionHolder, estimatorProto, Some(fitCmd.getParams))
val model = estimator.fit(dataset).asInstanceOf[Model[_]]
val id = mlCache.register(model)
proto.MlCommandResult
@@ -176,7 +177,8 @@ private[connect] object MLHandler extends Logging {
case proto.MlCommand.Write.TypeCase.OPERATOR =>
val writer = mlCommand.getWrite
if (writer.getOperator.getType == proto.MlOperator.OperatorType.ESTIMATOR) {
- val estimator = MLUtils.getEstimator(writer.getOperator, Some(writer.getParams))
+ val estimator =
+ MLUtils.getEstimator(sessionHolder, writer.getOperator, Some(writer.getParams))
estimator match {
case m: MLWritable => MLUtils.write(m, mlCommand.getWrite)
case other => throw MlUnsupportedException(s"Estimator $other is not writable")
@@ -195,7 +197,7 @@ private[connect] object MLHandler extends Logging {
val path = mlCommand.getRead.getPath
if (operator.getType == proto.MlOperator.OperatorType.MODEL) {
- val model = MLUtils.load(name, path).asInstanceOf[Model[_]]
+ val model = MLUtils.load(sessionHolder, name, path).asInstanceOf[Model[_]]
val id = mlCache.register(model)
proto.MlCommandResult
.newBuilder()
@@ -208,7 +210,7 @@ private[connect] object MLHandler extends Logging {
.build()
} else if (operator.getType == proto.MlOperator.OperatorType.ESTIMATOR) {
- val estimator = MLUtils.load(name, path).asInstanceOf[Estimator[_]]
+ val estimator = MLUtils.load(sessionHolder, name, path).asInstanceOf[Estimator[_]]
proto.MlCommandResult
.newBuilder()
.setOperatorInfo(
@@ -231,14 +233,14 @@ private[connect] object MLHandler extends Logging {
// Ml transform
case proto.MlRelation.MlTypeCase.TRANSFORM =>
relation.getTransform.getOperatorCase match {
- // transform for a new ML transformer
+ // transform with a new ML transformer
case proto.MlRelation.Transform.OperatorCase.TRANSFORMER =>
val transformProto = relation.getTransform
assert(
transformProto.getTransformer.getType ==
proto.MlOperator.OperatorType.TRANSFORMER)
val dataset = MLUtils.parseRelationProto(transformProto.getInput, sessionHolder)
- val transformer = MLUtils.getTransformer(transformProto)
+ val transformer = MLUtils.getTransformer(sessionHolder, transformProto)
transformer.transform(dataset)
// transform on a cached model
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLUtils.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLUtils.scala
index f3db36c7647d..2c006d86dfac 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLUtils.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLUtils.scala
@@ -17,7 +17,7 @@
package org.apache.spark.sql.connect.ml
-import java.util.ServiceLoader
+import java.util.{Optional, ServiceLoader}
import scala.collection.immutable.HashSet
import scala.jdk.CollectionConverters._
@@ -28,10 +28,11 @@ import org.apache.spark.connect.proto
import org.apache.spark.ml.{Estimator, Transformer}
import org.apache.spark.ml.linalg.{Matrices, Matrix, Vector, Vectors}
import org.apache.spark.ml.param.Params
-import org.apache.spark.ml.util.MLWritable
+import org.apache.spark.ml.util.{MLReadable, MLWritable}
import org.apache.spark.sql.{DataFrame, Dataset}
import org.apache.spark.sql.connect.common.LiteralValueProtoConverter
import org.apache.spark.sql.connect.planner.SparkConnectPlanner
+import org.apache.spark.sql.connect.plugin.SparkConnectPluginRegistry
import org.apache.spark.sql.connect.service.SessionHolder
import org.apache.spark.util.{SparkClassUtils, Utils}
@@ -41,7 +42,7 @@ private[ml] object MLUtils {
* Load the registered ML operators via ServiceLoader
*
* @param mlCls
- * the operator class
+ * the operator class that will be loaded
* @return
* a Map with name and class
*/
@@ -52,10 +53,6 @@ private[ml] object MLUtils {
providers.map(est => est.getClass.getName -> est.getClass).toMap
}
- private lazy val estimators = loadOperators(classOf[Estimator[_]])
-
- private lazy val transformers = loadOperators(classOf[Transformer])
-
def deserializeVector(vector: proto.Vector): Vector = {
if (vector.hasDense) {
val values = vector.getDense.getValueList.asScala.map(_.toDouble).toArray
@@ -228,9 +225,24 @@ private[ml] object MLUtils {
instance
}
+ /**
+ * Replace the operator with the value provided by the backend.
+ */
+ private def replaceOperator(sessionHolder: SessionHolder, name: String): String = {
+ SparkConnectPluginRegistry
+ .mlBackendRegistry(sessionHolder.session.sessionState.conf)
+ .view
+ .map(p => p.transform(name))
+ .find(_.isPresent) // First come, first served
+ .getOrElse(Optional.of(name))
+ .get()
+ }
+
/**
* Get the Estimator instance according to the proto information
*
+ * @param sessionHolder
+ * session holder to hold the Spark Connect session state
* @param operator
* MlOperator information
* @param params
@@ -238,29 +250,41 @@ private[ml] object MLUtils {
* @return
* the estimator
*/
- def getEstimator(operator: proto.MlOperator, params: Option[proto.MlParams]): Estimator[_] = {
- val name = operator.getName
+ def getEstimator(
+ sessionHolder: SessionHolder,
+ operator: proto.MlOperator,
+ params: Option[proto.MlParams]): Estimator[_] = {
+ val name = replaceOperator(sessionHolder, operator.getName)
val uid = operator.getUid
+
+ // Load the estimator by ServiceLoader everytime
+ val estimators = loadOperators(classOf[Estimator[_]])
getInstance[Estimator[_]](name, uid, estimators, params)
}
/**
* Get the transformer instance according to the transform proto
*
+ * @param sessionHolder
+ * session holder to hold the Spark Connect session state
* @param transformProto
* transform proto
* @return
* a transformer
*/
- def getTransformer(transformProto: proto.MlRelation.Transform): Transformer = {
- val name = transformProto.getTransformer.getName
+ def getTransformer(
+ sessionHolder: SessionHolder,
+ transformProto: proto.MlRelation.Transform): Transformer = {
+ val name = replaceOperator(sessionHolder, transformProto.getTransformer.getName)
val uid = transformProto.getTransformer.getUid
val params = transformProto.getParams
+ // Load the transformer by ServiceLoader everytime.
+ val transformers = loadOperators(classOf[Transformer])
getInstance[Transformer](name, uid, transformers, Some(params))
}
/**
- * Call "load: function on the ML operator given the operator name
+ * Call "load" function on the ML operator given the operator name
*
* @param className
* the ML operator name
@@ -269,8 +293,18 @@ private[ml] object MLUtils {
* @return
* the ML instance
*/
- def load(className: String, path: String): Object = {
- val loadedMethod = SparkClassUtils.classForName(className).getMethod("load", classOf[String])
+ def load(sessionHolder: SessionHolder, className: String, path: String): Object = {
+ val name = replaceOperator(sessionHolder, className)
+
+ // It's the companion object of the corresponding spark operators to load.
+ val objectCls = SparkClassUtils.classForName(name + "$")
+ val mlReadableClass = classOf[MLReadable[_]]
+ // Make sure it is implementing MLReadable
+ if (!mlReadableClass.isAssignableFrom(objectCls)) {
+ throw MlUnsupportedException(s"$name must implement MLReadable.")
+ }
+
+ val loadedMethod = SparkClassUtils.classForName(name).getMethod("load", classOf[String])
loadedMethod.invoke(null, path)
}
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/plugin/SparkConnectPluginRegistry.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/plugin/SparkConnectPluginRegistry.scala
index d6f5f01a5e07..fc4e7aed7aed 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/plugin/SparkConnectPluginRegistry.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/plugin/SparkConnectPluginRegistry.scala
@@ -21,6 +21,7 @@ import java.lang.reflect.InvocationTargetException
import org.apache.spark.{SparkEnv, SparkException}
import org.apache.spark.sql.connect.config.Connect
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.util.Utils
/**
@@ -64,6 +65,7 @@ object SparkConnectPluginRegistry {
def commandRegistry: Seq[CommandPlugin] = withInitialize {
commandRegistryCache
}
+ def mlBackendRegistry(conf: SQLConf): Seq[MLBackendPlugin] = loadMlBackendPlugins(conf)
private def withInitialize[T](f: => Seq[T]): Seq[T] = {
synchronized {
@@ -107,6 +109,10 @@ object SparkConnectPluginRegistry {
SparkEnv.get.conf.get(Connect.CONNECT_EXTENSIONS_COMMAND_CLASSES))
}
+ private[connect] def loadMlBackendPlugins(sqlConf: SQLConf): Seq[MLBackendPlugin] = {
+ createConfiguredPlugins(sqlConf.getConf(Connect.CONNECT_ML_BACKEND_CLASSES))
+ }
+
/**
* Exposed for testing only.
*/
diff --git a/sql/connect/server/src/test/resources/META-INF/services/org.apache.spark.ml.Estimator b/sql/connect/server/src/test/resources/META-INF/services/org.apache.spark.ml.Estimator
new file mode 100644
index 000000000000..37209d152b90
--- /dev/null
+++ b/sql/connect/server/src/test/resources/META-INF/services/org.apache.spark.ml.Estimator
@@ -0,0 +1,20 @@
+#
+# 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.
+#
+
+# Spark Connect ML uses ServiceLoader to find out the supported Spark Ml estimators.
+# So register the supported estimator here if you're trying to add a new one.
+org.apache.spark.sql.connect.ml.MyLogisticRegression
diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLBackendSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLBackendSuite.scala
new file mode 100644
index 000000000000..2e1d3d95c486
--- /dev/null
+++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLBackendSuite.scala
@@ -0,0 +1,193 @@
+/*
+ * 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.spark.sql.connect.ml
+
+import java.io.File
+
+import org.apache.spark.SparkEnv
+import org.apache.spark.connect.proto
+import org.apache.spark.ml.util.Identifiable
+import org.apache.spark.sql.connect.SparkConnectTestUtils
+import org.apache.spark.sql.connect.config.Connect
+import org.apache.spark.util.Utils
+
+class MLBackendSuite extends MLHelper {
+
+ def withSparkConf(pairs: (String, String)*)(f: => Unit): Unit = {
+ val conf = SparkEnv.get.conf
+ pairs.foreach { kv => conf.set(kv._1, kv._2) }
+ try f
+ finally {
+ pairs.foreach { kv => conf.remove(kv._1) }
+ }
+ }
+
+ private def getLogisticRegressionBuilder: proto.MlOperator.Builder = {
+ val name = "org.apache.spark.ml.classification.LogisticRegression"
+ proto.MlOperator
+ .newBuilder()
+ .setName(name)
+ .setUid(name)
+ .setType(proto.MlOperator.OperatorType.ESTIMATOR)
+ }
+
+ private def getMaxIterBuilder: proto.MlParams.Builder = {
+ proto.MlParams
+ .newBuilder()
+ .putParams(
+ "maxIter",
+ proto.Param
+ .newBuilder()
+ .setLiteral(
+ proto.Expression.Literal
+ .newBuilder()
+ .setInteger(2))
+ .build())
+ }
+
+ test("ML backend: estimator read/write") {
+ withSparkConf(
+ Connect.CONNECT_ML_BACKEND_CLASSES.key ->
+ "org.apache.spark.sql.connect.ml.MyMlBackend") {
+ val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark)
+
+ // read/write
+ val tempDir = Utils.createTempDir(namePrefix = this.getClass.getName)
+ try {
+ val path = new File(tempDir, Identifiable.randomUID("LogisticRegression")).getPath
+ val writeCmd = proto.MlCommand
+ .newBuilder()
+ .setWrite(
+ proto.MlCommand.Write
+ .newBuilder()
+ .setOperator(getLogisticRegressionBuilder)
+ .setParams(getMaxIterBuilder)
+ .setPath(path)
+ .setShouldOverwrite(true))
+ .build()
+ MLHandler.handleMlCommand(sessionHolder, writeCmd)
+
+ val readCmd = proto.MlCommand
+ .newBuilder()
+ .setRead(
+ proto.MlCommand.Read
+ .newBuilder()
+ .setOperator(getLogisticRegressionBuilder)
+ .setPath(path))
+ .build()
+
+ val ret = MLHandler.handleMlCommand(sessionHolder, readCmd)
+ assert(ret.getOperatorInfo.getParams.getParamsMap.containsKey("fakeParam"))
+ assert(ret.getOperatorInfo.getParams.getParamsMap.containsKey("maxIter"))
+ assert(
+ ret.getOperatorInfo.getParams.getParamsMap.get("maxIter").getLiteral.getInteger
+ == 2)
+ assert(
+ ret.getOperatorInfo.getParams.getParamsMap.get("fakeParam").getLiteral.getInteger
+ == 101010)
+ } finally {
+ Utils.deleteRecursively(tempDir)
+ }
+ }
+ }
+
+ test("ML backend: model read/write") {
+ withSparkConf(
+ Connect.CONNECT_ML_BACKEND_CLASSES.key ->
+ "org.apache.spark.sql.connect.ml.MyMlBackend") {
+ val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark)
+
+ val tempDir = Utils.createTempDir(namePrefix = this.getClass.getName)
+ try {
+ // Train a model
+ val fitCommand = proto.MlCommand
+ .newBuilder()
+ .setFit(
+ proto.MlCommand.Fit
+ .newBuilder()
+ .setDataset(createLocalRelationProto)
+ .setEstimator(getLogisticRegressionBuilder)
+ .setParams(getMaxIterBuilder))
+ .build()
+ val fitRet = MLHandler.handleMlCommand(sessionHolder, fitCommand)
+ val modelId = fitRet.getOperatorInfo.getObjRef.getId
+
+ // Write a model
+ val path = new File(tempDir, Identifiable.randomUID("LogisticRegression")).getPath
+ val writeCmd = proto.MlCommand
+ .newBuilder()
+ .setWrite(
+ proto.MlCommand.Write
+ .newBuilder()
+ .setObjRef(proto.ObjectRef.newBuilder().setId(modelId))
+ .setPath(path)
+ .setShouldOverwrite(true))
+ .build()
+ MLHandler.handleMlCommand(sessionHolder, writeCmd)
+
+ // read a model
+ val readCmd = proto.MlCommand
+ .newBuilder()
+ .setRead(
+ proto.MlCommand.Read
+ .newBuilder()
+ .setOperator(proto.MlOperator
+ .newBuilder()
+ .setName("org.apache.spark.ml.classification.LogisticRegressionModel")
+ .setType(proto.MlOperator.OperatorType.MODEL))
+ .setPath(path))
+ .build()
+
+ val ret = MLHandler.handleMlCommand(sessionHolder, readCmd)
+ assert(ret.getOperatorInfo.getParams.getParamsMap.containsKey("fakeParam"))
+ assert(ret.getOperatorInfo.getParams.getParamsMap.containsKey("maxIter"))
+ assert(
+ ret.getOperatorInfo.getParams.getParamsMap.get("maxIter").getLiteral.getInteger
+ == 2)
+ assert(
+ ret.getOperatorInfo.getParams.getParamsMap.get("fakeParam").getLiteral.getInteger
+ == 101010)
+ } finally {
+ Utils.deleteRecursively(tempDir)
+ }
+ }
+ }
+
+ test("ML backend") {
+ withSparkConf(
+ Connect.CONNECT_ML_BACKEND_CLASSES.key ->
+ "org.apache.spark.sql.connect.ml.MyMlBackend") {
+ val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark)
+ val fitCommand = proto.MlCommand
+ .newBuilder()
+ .setFit(
+ proto.MlCommand.Fit
+ .newBuilder()
+ .setDataset(createLocalRelationProto)
+ .setEstimator(getLogisticRegressionBuilder)
+ .setParams(getMaxIterBuilder))
+ .build()
+ val fitResult = MLHandler.handleMlCommand(sessionHolder, fitCommand)
+ val modelId = fitResult.getOperatorInfo.getObjRef.getId
+ assert(sessionHolder.mlCache.get(modelId).isInstanceOf[MyLogisticRegressionModel])
+ val model = sessionHolder.mlCache.get(modelId).asInstanceOf[MyLogisticRegressionModel]
+ assert(model.intercept == 3.5f)
+ assert(model.coefficients == 4.6f)
+ }
+ }
+}
diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLHelper.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLHelper.scala
new file mode 100644
index 000000000000..c778603eeece
--- /dev/null
+++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLHelper.scala
@@ -0,0 +1,138 @@
+/*
+ * 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.spark.sql.connect.ml
+
+import java.util.Optional
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.connect.proto
+import org.apache.spark.ml.{Estimator, Model}
+import org.apache.spark.ml.linalg.{Vectors, VectorUDT}
+import org.apache.spark.ml.param.{IntParam, ParamMap, Params}
+import org.apache.spark.ml.param.shared.HasMaxIter
+import org.apache.spark.ml.util.{DefaultParamsReadable, DefaultParamsWritable, Identifiable, MLReadable, MLReader}
+import org.apache.spark.sql.{DataFrame, Dataset}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.UnsafeProjection
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.connect.planner.SparkConnectPlanTest
+import org.apache.spark.sql.connect.plugin.MLBackendPlugin
+import org.apache.spark.sql.types.{FloatType, Metadata, StructField, StructType}
+
+trait MLHelper extends SparkFunSuite with SparkConnectPlanTest {
+
+ def createLocalRelationProto: proto.Relation = {
+ val udt = new VectorUDT()
+ val rows = Seq(
+ InternalRow(1.0f, udt.serialize(Vectors.dense(Array(1.0, 2.0)))),
+ InternalRow(1.0f, udt.serialize(Vectors.dense(Array(2.0, -1.0)))),
+ InternalRow(0.0f, udt.serialize(Vectors.dense(Array(-3.0, -2.0)))),
+ InternalRow(0.0f, udt.serialize(Vectors.dense(Array(-1.0, -2.0)))))
+
+ val schema = StructType(
+ Seq(
+ StructField("label", FloatType),
+ StructField("features", new VectorUDT(), false, Metadata.empty)))
+
+ val inputRows = rows.map { row =>
+ val proj = UnsafeProjection.create(schema)
+ proj(row).copy()
+ }
+ createLocalRelationProto(DataTypeUtils.toAttributes(schema), inputRows, "UTC", Some(schema))
+ }
+}
+
+class MyMlBackend extends MLBackendPlugin {
+
+ override def transform(mlName: String): Optional[String] = {
+ mlName match {
+ case "org.apache.spark.ml.classification.LogisticRegression" =>
+ Optional.of("org.apache.spark.sql.connect.ml.MyLogisticRegression")
+ case "org.apache.spark.ml.classification.LogisticRegressionModel" =>
+ Optional.of("org.apache.spark.sql.connect.ml.MyLogisticRegressionModel")
+ case _ => Optional.empty()
+ }
+ }
+}
+
+trait HasFakedParam extends Params {
+ final val fakeParam: IntParam = new IntParam(this, "fakeParam", "faked parameter")
+}
+
+class MyLogisticRegressionModel(
+ override val uid: String,
+ val intercept: Float,
+ val coefficients: Float)
+ extends Model[MyLogisticRegressionModel]
+ with HasMaxIter
+ with HasFakedParam
+ with DefaultParamsWritable {
+
+ def setFakeParam(v: Int): this.type = set(fakeParam, v)
+
+ def setMaxIter(v: Int): this.type = set(maxIter, v)
+
+ override def copy(extra: ParamMap): MyLogisticRegressionModel = {
+ copyValues(new MyLogisticRegressionModel(uid, intercept, coefficients), extra)
+ }
+
+ override def transform(dataset: Dataset[_]): DataFrame = {
+ dataset.toDF()
+ }
+
+ override def transformSchema(schema: StructType): StructType = schema
+}
+
+object MyLogisticRegressionModel extends MLReadable[MyLogisticRegressionModel] {
+
+ // No need to load from file.
+ override def read: MLReader[MyLogisticRegressionModel] =
+ (_: String) => {
+ new MyLogisticRegressionModel("MyLogisticRegressionModel", 3.5f, 4.6f)
+ .setMaxIter(2)
+ .setFakeParam(101010)
+ }
+}
+
+class MyLogisticRegression(override val uid: String)
+ extends Estimator[MyLogisticRegressionModel]
+ with HasMaxIter
+ with HasFakedParam
+ with DefaultParamsWritable {
+ set(fakeParam, 101010)
+
+ def this() = this(Identifiable.randomUID("MyLogisticRegression"))
+
+ override def fit(dataset: Dataset[_]): MyLogisticRegressionModel = {
+ copyValues(new MyLogisticRegressionModel(uid, 3.5f, 4.6f))
+ }
+
+ override def copy(extra: ParamMap): MyLogisticRegression = {
+ defaultCopy(extra).asInstanceOf[MyLogisticRegression]
+ }
+
+ override def transformSchema(schema: StructType): StructType = schema
+}
+
+object MyLogisticRegression extends DefaultParamsReadable[MyLogisticRegression] {
+ override def load(path: String): MyLogisticRegression = super.load(path)
+}
+
+object NotImplementingMLReadble {
+ def load(path: String): Unit = {}
+}
diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLSuite.scala
index bea7072b2034..280ced32a3f8 100644
--- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLSuite.scala
+++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLSuite.scala
@@ -19,18 +19,13 @@ package org.apache.spark.sql.connect.ml
import java.io.File
-import org.apache.spark.SparkFunSuite
import org.apache.spark.connect.proto
import org.apache.spark.ml.classification.LogisticRegressionModel
-import org.apache.spark.ml.linalg.{Vectors, VectorUDT}
+import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.ml.param._
import org.apache.spark.ml.util.Identifiable
-import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.UnsafeProjection
-import org.apache.spark.sql.catalyst.types.DataTypeUtils
import org.apache.spark.sql.connect.SparkConnectTestUtils
-import org.apache.spark.sql.connect.planner.SparkConnectPlanTest
-import org.apache.spark.sql.types.{FloatType, Metadata, StructField, StructType}
+import org.apache.spark.sql.connect.service.SessionHolder
import org.apache.spark.util.Utils
trait FakeArrayParams extends Params {
@@ -71,27 +66,7 @@ class FakedML(override val uid: String) extends FakeArrayParams {
override def copy(extra: ParamMap): Params = this
}
-class MLSuite extends SparkFunSuite with SparkConnectPlanTest {
-
- def createLocalRelationProto: proto.Relation = {
- val udt = new VectorUDT()
- val rows = Seq(
- InternalRow(1.0f, udt.serialize(Vectors.dense(Array(1.0, 2.0)))),
- InternalRow(1.0f, udt.serialize(Vectors.dense(Array(2.0, -1.0)))),
- InternalRow(0.0f, udt.serialize(Vectors.dense(Array(-3.0, -2.0)))),
- InternalRow(0.0f, udt.serialize(Vectors.dense(Array(-1.0, -2.0)))))
-
- val schema = StructType(
- Seq(
- StructField("label", FloatType),
- StructField("features", new VectorUDT(), false, Metadata.empty)))
-
- val inputRows = rows.map { row =>
- val proj = UnsafeProjection.create(schema)
- proj(row).copy()
- }
- createLocalRelationProto(DataTypeUtils.toAttributes(schema), inputRows, "UTC", Some(schema))
- }
+class MLSuite extends MLHelper {
test("reconcileParam") {
val fakedML = new FakedML
@@ -189,6 +164,46 @@ class MLSuite extends SparkFunSuite with SparkConnectPlanTest {
assert(fakedML.getDouble === 1.0)
}
+ def fetchCommand(modelId: String, method: String): proto.MlCommand = {
+ proto.MlCommand
+ .newBuilder()
+ .setFetch(
+ proto.Fetch
+ .newBuilder()
+ .setObjRef(proto.ObjectRef.newBuilder().setId(modelId))
+ .addMethods(proto.Fetch.Method.newBuilder().setMethod(method)))
+ .build()
+ }
+
+ def trainLogisticRegressionModel(sessionHolder: SessionHolder): String = {
+ val fitCommand = proto.MlCommand
+ .newBuilder()
+ .setFit(
+ proto.MlCommand.Fit
+ .newBuilder()
+ .setDataset(createLocalRelationProto)
+ .setEstimator(
+ proto.MlOperator
+ .newBuilder()
+ .setName("org.apache.spark.ml.classification.LogisticRegression")
+ .setUid("LogisticRegression")
+ .setType(proto.MlOperator.OperatorType.ESTIMATOR))
+ .setParams(
+ proto.MlParams
+ .newBuilder()
+ .putParams(
+ "maxIter",
+ proto.Param
+ .newBuilder()
+ .setLiteral(proto.Expression.Literal
+ .newBuilder()
+ .setInteger(2))
+ .build())))
+ .build()
+ val fitResult = MLHandler.handleMlCommand(sessionHolder, fitCommand)
+ fitResult.getOperatorInfo.getObjRef.getId
+ }
+
test("LogisticRegression works") {
val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark)
@@ -201,40 +216,19 @@ class MLSuite extends SparkFunSuite with SparkConnectPlanTest {
assert(lrModel.getMaxIter === 2)
// Fetch double attribute
- val interceptCommand = proto.MlCommand
- .newBuilder()
- .setFetch(
- proto.Fetch
- .newBuilder()
- .setObjRef(proto.ObjectRef.newBuilder().setId(modelId))
- .addMethods(proto.Fetch.Method.newBuilder().setMethod("intercept")))
- .build()
+ val interceptCommand = fetchCommand(modelId, "intercept")
val interceptResult = MLHandler.handleMlCommand(sessionHolder, interceptCommand)
assert(interceptResult.getParam.getLiteral.getDouble === lrModel.intercept)
// Fetch Vector attribute
- val coefficientsCommand = proto.MlCommand
- .newBuilder()
- .setFetch(
- proto.Fetch
- .newBuilder()
- .setObjRef(proto.ObjectRef.newBuilder().setId(modelId))
- .addMethods(proto.Fetch.Method.newBuilder().setMethod("coefficients")))
- .build()
+ val coefficientsCommand = fetchCommand(modelId, "coefficients")
val coefficientsResult = MLHandler.handleMlCommand(sessionHolder, coefficientsCommand)
val deserializedCoefficients =
MLUtils.deserializeVector(coefficientsResult.getParam.getVector)
assert(deserializedCoefficients === lrModel.coefficients)
// Fetch Matrix attribute
- val coefficientsMatrixCommand = proto.MlCommand
- .newBuilder()
- .setFetch(
- proto.Fetch
- .newBuilder()
- .setObjRef(proto.ObjectRef.newBuilder().setId(modelId))
- .addMethods(proto.Fetch.Method.newBuilder().setMethod("coefficientMatrix")))
- .build()
+ val coefficientsMatrixCommand = fetchCommand(modelId, "coefficientMatrix")
val coefficientsMatrixResult =
MLHandler.handleMlCommand(sessionHolder, coefficientsMatrixCommand)
val deserializedCoefficientsMatrix =
@@ -300,32 +294,7 @@ class MLSuite extends SparkFunSuite with SparkConnectPlanTest {
}
try {
- val fitCommand = proto.MlCommand
- .newBuilder()
- .setFit(
- proto.MlCommand.Fit
- .newBuilder()
- .setDataset(createLocalRelationProto)
- .setEstimator(
- proto.MlOperator
- .newBuilder()
- .setName("org.apache.spark.ml.classification.LogisticRegression")
- .setUid("LogisticRegression")
- .setType(proto.MlOperator.OperatorType.ESTIMATOR))
- .setParams(
- proto.MlParams
- .newBuilder()
- .putParams(
- "maxIter",
- proto.Param
- .newBuilder()
- .setLiteral(proto.Expression.Literal
- .newBuilder()
- .setInteger(2))
- .build())))
- .build()
- val fitResult = MLHandler.handleMlCommand(sessionHolder, fitCommand)
- val modelId = fitResult.getOperatorInfo.getObjRef.getId
+ val modelId = trainLogisticRegressionModel(sessionHolder)
verifyModel(modelId, true)
@@ -387,4 +356,35 @@ class MLSuite extends SparkFunSuite with SparkConnectPlanTest {
MLHandler.handleMlCommand(sessionHolder, command)
}
}
+
+ test("access the attribute which is not in allowed list") {
+ val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark)
+ val modelId = trainLogisticRegressionModel(sessionHolder)
+
+ val fakeAttributeCmd = fetchCommand(modelId, "notExistingAttribute")
+ intercept[MLAttributeNotAllowedException] {
+ MLHandler.handleMlCommand(sessionHolder, fakeAttributeCmd)
+ }
+ }
+
+ test("ML operator must implement MLReadable for loading") {
+ val thrown = intercept[MlUnsupportedException] {
+ val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark)
+ val readCmd = proto.MlCommand
+ .newBuilder()
+ .setRead(
+ proto.MlCommand.Read
+ .newBuilder()
+ .setOperator(proto.MlOperator
+ .newBuilder()
+ .setName("org.apache.spark.sql.connect.ml.NotImplementingMLReadble")
+ .setType(proto.MlOperator.OperatorType.ESTIMATOR))
+ .setPath("/tmp/fake"))
+ .build()
+ MLHandler.handleMlCommand(sessionHolder, readCmd)
+ }
+ assert(
+ thrown.message.contains("org.apache.spark.sql.connect.ml.NotImplementingMLReadble " +
+ "must implement MLReadable"))
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org