[SPARK-52470][ML][CONNECT] Support model summary offloading

### What changes were proposed in this pull request?

This PR makes Spark Connect ML supporting model summary offloading.

Model summary offloading is hard to support because it contains a Spark dataset which can't be easily serialized in Spark driver (NOTE: we can't java serializer to serialize the Spark dataset logical plan otherwise it is a RCE vulnerability),

to address the issue, when saving Summary to disk, it only saves the necessary data fields,
when loading Summary back, the client needs to send the dataset to Spark driver again,
to achieve it, 2 new proto messages are introduced:

1. `CreateSummary` in `MlCommand`
```
  // This is for re-creating the model summary when the model summary is lost
  // (model summary is lost when the model is offloaded and then loaded back)
  message CreateSummary {
    ObjectRef model_ref = 1;
    Relation dataset = 2;
  }
```

2: `model_summary_dataset` in `MlRelation`

```
  // (Optional) the dataset for restoring the model summary
  optional Relation model_summary_dataset = 3;
```

### Why are the changes needed?

Support model summary offloading.
Without this, the model summary will be evicted from Spark driver memory after default 15min timeout, results in `model.summary` API unavailability.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Unit tests.

### Was this patch authored or co-authored using generative AI tooling?

No.

Closes #51187 from WeichenXu123/SPARK-52470.

Authored-by: Weichen Xu <[email protected]>
Signed-off-by: Weichen Xu <[email protected]>

Author: WeichenXu123

common/utils/src/main/resources/error/error-conditions.json

@@ -853,6 +853,11 @@
           "Please fit or load a model smaller than <modelMaxSize> bytes."
         ]
       },
+      "MODEL_SUMMARY_LOST" : {
+        "message" : [
+          "The model <objectName> summary is lost because the cached model is offloaded."
+        ]
+      },
       "UNSUPPORTED_EXCEPTION" : {
         "message" : [
           "<message>"

mllib/src/main/scala/org/apache/spark/ml/classification/FMClassifier.scala

@@ -224,17 +224,8 @@ class FMClassifier @Since("3.0.0") (
     factors: Matrix,
     objectiveHistory: Array[Double]): FMClassificationModel = {
     val model = copyValues(new FMClassificationModel(uid, intercept, linear, factors))
-    val weightColName = if (!isDefined(weightCol)) "weightCol" else $(weightCol)
-
-    val (summaryModel, probabilityColName, predictionColName) = model.findSummaryModel()
-    val summary = new FMClassificationTrainingSummaryImpl(
-      summaryModel.transform(dataset),
-      probabilityColName,
-      predictionColName,
-      $(labelCol),
-      weightColName,
-      objectiveHistory)
-    model.setSummary(Some(summary))
+    model.createSummary(dataset, objectiveHistory)
+    model
   }
 
   @Since("3.0.0")
@@ -343,6 +334,42 @@ class FMClassificationModel private[classification] (
       s"uid=${super.toString}, numClasses=$numClasses, numFeatures=$numFeatures, " +
       s"factorSize=${$(factorSize)}, fitLinear=${$(fitLinear)}, fitIntercept=${$(fitIntercept)}"
   }
+
+  private[spark] def createSummary(
+    dataset: Dataset[_], objectiveHistory: Array[Double]
+  ): Unit = {
+    val weightColName = if (!isDefined(weightCol)) "weightCol" else $(weightCol)
+
+    val (summaryModel, probabilityColName, predictionColName) = findSummaryModel()
+    val summary = new FMClassificationTrainingSummaryImpl(
+      summaryModel.transform(dataset),
+      probabilityColName,
+      predictionColName,
+      $(labelCol),
+      weightColName,
+      objectiveHistory)
+    setSummary(Some(summary))
+  }
+
+  override private[spark] def saveSummary(path: String): Unit = {
+    ReadWriteUtils.saveObjectToLocal[Tuple1[Array[Double]]](
+      path, Tuple1(summary.objectiveHistory),
+      (data, dos) => {
+        ReadWriteUtils.serializeDoubleArray(data._1, dos)
+      }
+    )
+  }
+
+  override private[spark] def loadSummary(path: String, dataset: DataFrame): Unit = {
+    val Tuple1(objectiveHistory: Array[Double])
+    = ReadWriteUtils.loadObjectFromLocal[Tuple1[Array[Double]]](
+      path,
+      dis => {
+        Tuple1(ReadWriteUtils.deserializeDoubleArray(dis))
+      }
+    )
+    createSummary(dataset, objectiveHistory)
+  }
 }
 
 @Since("3.0.0")

mllib/src/main/scala/org/apache/spark/ml/classification/LinearSVC.scala

@@ -277,17 +277,8 @@ class LinearSVC @Since("2.2.0") (
       intercept: Double,
       objectiveHistory: Array[Double]): LinearSVCModel = {
     val model = copyValues(new LinearSVCModel(uid, coefficients, intercept))
-    val weightColName = if (!isDefined(weightCol)) "weightCol" else $(weightCol)
-
-    val (summaryModel, rawPredictionColName, predictionColName) = model.findSummaryModel()
-    val summary = new LinearSVCTrainingSummaryImpl(
-      summaryModel.transform(dataset),
-      rawPredictionColName,
-      predictionColName,
-      $(labelCol),
-      weightColName,
-      objectiveHistory)
-    model.setSummary(Some(summary))
+    model.createSummary(dataset, objectiveHistory)
+    model
   }
 
   private def trainImpl(
@@ -445,6 +436,42 @@ class LinearSVCModel private[classification] (
   override def toString: String = {
     s"LinearSVCModel: uid=$uid, numClasses=$numClasses, numFeatures=$numFeatures"
   }
+
+  private[spark] def createSummary(
+    dataset: Dataset[_], objectiveHistory: Array[Double]
+  ): Unit = {
+    val weightColName = if (!isDefined(weightCol)) "weightCol" else $(weightCol)
+
+    val (summaryModel, rawPredictionColName, predictionColName) = findSummaryModel()
+    val summary = new LinearSVCTrainingSummaryImpl(
+      summaryModel.transform(dataset),
+      rawPredictionColName,
+      predictionColName,
+      $(labelCol),
+      weightColName,
+      objectiveHistory)
+    setSummary(Some(summary))
+  }
+
+  override private[spark] def saveSummary(path: String): Unit = {
+    ReadWriteUtils.saveObjectToLocal[Tuple1[Array[Double]]](
+      path, Tuple1(summary.objectiveHistory),
+      (data, dos) => {
+        ReadWriteUtils.serializeDoubleArray(data._1, dos)
+      }
+    )
+  }
+
+  override private[spark] def loadSummary(path: String, dataset: DataFrame): Unit = {
+    val Tuple1(objectiveHistory: Array[Double])
+    = ReadWriteUtils.loadObjectFromLocal[Tuple1[Array[Double]]](
+      path,
+      dis => {
+        Tuple1(ReadWriteUtils.deserializeDoubleArray(dis))
+      }
+    )
+    createSummary(dataset, objectiveHistory)
+  }
 }
 
 @Since("2.2.0")

mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala

@@ -718,29 +718,8 @@ class LogisticRegression @Since("1.2.0") (
       objectiveHistory: Array[Double]): LogisticRegressionModel = {
     val model = copyValues(new LogisticRegressionModel(uid, coefficientMatrix, interceptVector,
       numClasses, checkMultinomial(numClasses)))
-    val weightColName = if (!isDefined(weightCol)) "weightCol" else $(weightCol)
-
-    val (summaryModel, probabilityColName, predictionColName) = model.findSummaryModel()
-    val logRegSummary = if (numClasses <= 2) {
-      new BinaryLogisticRegressionTrainingSummaryImpl(
-        summaryModel.transform(dataset),
-        probabilityColName,
-        predictionColName,
-        $(labelCol),
-        $(featuresCol),
-        weightColName,
-        objectiveHistory)
-    } else {
-      new LogisticRegressionTrainingSummaryImpl(
-        summaryModel.transform(dataset),
-        probabilityColName,
-        predictionColName,
-        $(labelCol),
-        $(featuresCol),
-        weightColName,
-        objectiveHistory)
-    }
-    model.setSummary(Some(logRegSummary))
+    model.createSummary(dataset, objectiveHistory)
+    model
   }
 
   private def createBounds(
@@ -1323,6 +1302,54 @@ class LogisticRegressionModel private[spark] (
   override def toString: String = {
     s"LogisticRegressionModel: uid=$uid, numClasses=$numClasses, numFeatures=$numFeatures"
   }
+
+  private[spark] def createSummary(
+    dataset: Dataset[_], objectiveHistory: Array[Double]
+  ): Unit = {
+    val weightColName = if (!isDefined(weightCol)) "weightCol" else $(weightCol)
+
+    val (summaryModel, probabilityColName, predictionColName) = findSummaryModel()
+    val logRegSummary = if (numClasses <= 2) {
+      new BinaryLogisticRegressionTrainingSummaryImpl(
+        summaryModel.transform(dataset),
+        probabilityColName,
+        predictionColName,
+        $(labelCol),
+        $(featuresCol),
+        weightColName,
+        objectiveHistory)
+    } else {
+      new LogisticRegressionTrainingSummaryImpl(
+        summaryModel.transform(dataset),
+        probabilityColName,
+        predictionColName,
+        $(labelCol),
+        $(featuresCol),
+        weightColName,
+        objectiveHistory)
+    }
+    setSummary(Some(logRegSummary))
+  }
+
+  override private[spark] def saveSummary(path: String): Unit = {
+    ReadWriteUtils.saveObjectToLocal[Tuple1[Array[Double]]](
+      path, Tuple1(summary.objectiveHistory),
+      (data, dos) => {
+        ReadWriteUtils.serializeDoubleArray(data._1, dos)
+      }
+    )
+  }
+
+  override private[spark] def loadSummary(path: String, dataset: DataFrame): Unit = {
+    val Tuple1(objectiveHistory: Array[Double])
+    = ReadWriteUtils.loadObjectFromLocal[Tuple1[Array[Double]]](
+      path,
+      dis => {
+        Tuple1(ReadWriteUtils.deserializeDoubleArray(dis))
+      }
+    )
+    createSummary(dataset, objectiveHistory)
+  }
 }
 
 @Since("1.6.0")

mllib/src/main/scala/org/apache/spark/ml/classification/MultilayerPerceptronClassifier.scala

@@ -251,14 +251,8 @@ class MultilayerPerceptronClassifier @Since("1.5.0") (
       objectiveHistory: Array[Double]): MultilayerPerceptronClassificationModel = {
     val model = copyValues(new MultilayerPerceptronClassificationModel(uid, weights))
 
-    val (summaryModel, _, predictionColName) = model.findSummaryModel()
-    val summary = new MultilayerPerceptronClassificationTrainingSummaryImpl(
-      summaryModel.transform(dataset),
-      predictionColName,
-      $(labelCol),
-      "",
-      objectiveHistory)
-    model.setSummary(Some(summary))
+    model.createSummary(dataset, objectiveHistory)
+    model
   }
 }
 
@@ -365,6 +359,39 @@ class MultilayerPerceptronClassificationModel private[ml] (
     s"MultilayerPerceptronClassificationModel: uid=$uid, numLayers=${$(layers).length}, " +
       s"numClasses=$numClasses, numFeatures=$numFeatures"
   }
+
+  private[spark] def createSummary(
+    dataset: Dataset[_], objectiveHistory: Array[Double]
+  ): Unit = {
+    val (summaryModel, _, predictionColName) = findSummaryModel()
+    val summary = new MultilayerPerceptronClassificationTrainingSummaryImpl(
+      summaryModel.transform(dataset),
+      predictionColName,
+      $(labelCol),
+      "",
+      objectiveHistory)
+    setSummary(Some(summary))
+  }
+
+  override private[spark] def saveSummary(path: String): Unit = {
+    ReadWriteUtils.saveObjectToLocal[Tuple1[Array[Double]]](
+      path, Tuple1(summary.objectiveHistory),
+      (data, dos) => {
+        ReadWriteUtils.serializeDoubleArray(data._1, dos)
+      }
+    )
+  }
+
+  override private[spark] def loadSummary(path: String, dataset: DataFrame): Unit = {
+    val Tuple1(objectiveHistory: Array[Double])
+    = ReadWriteUtils.loadObjectFromLocal[Tuple1[Array[Double]]](
+      path,
+      dis => {
+        Tuple1(ReadWriteUtils.deserializeDoubleArray(dis))
+      }
+    )
+    createSummary(dataset, objectiveHistory)
+  }
 }
 
 @Since("2.0.0")

mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala

@@ -182,26 +182,8 @@ class RandomForestClassifier @Since("1.4.0") (
       numFeatures: Int,
       numClasses: Int): RandomForestClassificationModel = {
     val model = copyValues(new RandomForestClassificationModel(uid, trees, numFeatures, numClasses))
-    val weightColName = if (!isDefined(weightCol)) "weightCol" else $(weightCol)
-
-    val (summaryModel, probabilityColName, predictionColName) = model.findSummaryModel()
-    val rfSummary = if (numClasses <= 2) {
-      new BinaryRandomForestClassificationTrainingSummaryImpl(
-        summaryModel.transform(dataset),
-        probabilityColName,
-        predictionColName,
-        $(labelCol),
-        weightColName,
-        Array(0.0))
-    } else {
-      new RandomForestClassificationTrainingSummaryImpl(
-        summaryModel.transform(dataset),
-        predictionColName,
-        $(labelCol),
-        weightColName,
-        Array(0.0))
-    }
-    model.setSummary(Some(rfSummary))
+    model.createSummary(dataset)
+    model
   }
 
   @Since("1.4.1")
@@ -393,6 +375,35 @@ class RandomForestClassificationModel private[ml] (
   @Since("2.0.0")
   override def write: MLWriter =
     new RandomForestClassificationModel.RandomForestClassificationModelWriter(this)
+
+  private[spark] def createSummary(dataset: Dataset[_]): Unit = {
+    val weightColName = if (!isDefined(weightCol)) "weightCol" else $(weightCol)
+
+    val (summaryModel, probabilityColName, predictionColName) = findSummaryModel()
+    val rfSummary = if (numClasses <= 2) {
+      new BinaryRandomForestClassificationTrainingSummaryImpl(
+        summaryModel.transform(dataset),
+        probabilityColName,
+        predictionColName,
+        $(labelCol),
+        weightColName,
+        Array(0.0))
+    } else {
+      new RandomForestClassificationTrainingSummaryImpl(
+        summaryModel.transform(dataset),
+        predictionColName,
+        $(labelCol),
+        weightColName,
+        Array(0.0))
+    }
+    setSummary(Some(rfSummary))
+  }
+
+  override private[spark] def saveSummary(path: String): Unit = {}
+
+  override private[spark] def loadSummary(path: String, dataset: DataFrame): Unit = {
+    createSummary(dataset)
+  }
 }
 
 @Since("2.0.0")

mllib/src/main/scala/org/apache/spark/ml/clustering/BisectingKMeans.scala

@@ -180,6 +180,9 @@ class BisectingKMeansModel private[ml] (
   override def summary: BisectingKMeansSummary = super.summary
 
   override def estimatedSize: Long = SizeEstimator.estimate(parentModel)
+
+  // BisectingKMeans model hasn't supported offloading, so put an empty `saveSummary` here for now
+  override private[spark] def saveSummary(path: String): Unit = {}
 }
 
 object BisectingKMeansModel extends MLReadable[BisectingKMeansModel] {