[SPARK-57870][ML][PYTHON] Apply HasIntermediateStorageLevel to tree ensemble estimators (GBT, RandomForest)

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

@@ -36,6 +36,7 @@ import org.apache.spark.mllib.tree.model.{GradientBoostedTreesModel => OldGBTMod
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types.StructType
+import org.apache.spark.storage.StorageLevel
 
 /**
  * Gradient-Boosted Trees (GBTs) (http://en.wikipedia.org/wiki/Gradient_boosting)
@@ -166,6 +167,10 @@ class GBTClassifier @Since("1.4.0") (
   @Since("3.0.0")
   def setWeightCol(value: String): this.type = set(weightCol, value)
 
+  /** @group expertSetParam */
+  @Since("5.0.0")
+  def setIntermediateStorageLevel(value: String): this.type = set(intermediateStorageLevel, value)
+
   override protected def train(
       dataset: Dataset[_]): GBTClassificationModel = instrumented { instr =>
     val withValidation = isDefined(validationIndicatorCol) && $(validationIndicatorCol).nonEmpty
@@ -188,17 +193,19 @@ class GBTClassifier @Since("1.4.0") (
     instr.logParams(this, labelCol, weightCol, featuresCol, predictionCol, leafCol,
       impurity, lossType, maxDepth, maxBins, maxIter, maxMemoryInMB, minInfoGain,
       minInstancesPerNode, minWeightFractionPerNode, seed, stepSize, subsamplingRate, cacheNodeIds,
-      checkpointInterval, featureSubsetStrategy, validationIndicatorCol, validationTol, thresholds)
+      checkpointInterval, featureSubsetStrategy, validationIndicatorCol, validationTol, thresholds,
+      intermediateStorageLevel)
     instr.logNumClasses(numClasses)
 
     val categoricalFeatures = MetadataUtils.getCategoricalFeatures(dataset.schema($(featuresCol)))
     val boostingStrategy = super.getOldBoostingStrategy(categoricalFeatures, OldAlgo.Classification)
+    val storageLevel = StorageLevel.fromString($(intermediateStorageLevel))
     val (baseLearners, learnerWeights) = if (withValidation) {
       GradientBoostedTrees.runWithValidation(trainDataset, validationDataset, boostingStrategy,
-        $(seed), $(featureSubsetStrategy), Some(instr))
+        $(seed), $(featureSubsetStrategy), Some(instr), storageLevel)
     } else {
       GradientBoostedTrees.run(trainDataset, boostingStrategy, $(seed), $(featureSubsetStrategy),
-        Some(instr))
+        Some(instr), storageLevel)
     }
     baseLearners.foreach(copyValues(_))
 

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

@@ -36,6 +36,7 @@ import org.apache.spark.mllib.tree.model.{RandomForestModel => OldRandomForestMo
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions.{col, udf}
 import org.apache.spark.sql.types.StructType
+import org.apache.spark.storage.StorageLevel
 
 /**
  * <a href="http://en.wikipedia.org/wiki/Random_forest">Random Forest</a> learning algorithm for
@@ -139,6 +140,10 @@ class RandomForestClassifier @Since("1.4.0") (
   @Since("3.0.0")
   def setWeightCol(value: String): this.type = set(weightCol, value)
 
+  /** @group expertSetParam */
+  @Since("5.0.0")
+  def setIntermediateStorageLevel(value: String): this.type = set(intermediateStorageLevel, value)
+
   override protected def train(
       dataset: Dataset[_]): RandomForestClassificationModel = instrumented { instr =>
     instr.logPipelineStage(this)
@@ -168,10 +173,12 @@ class RandomForestClassifier @Since("1.4.0") (
     instr.logParams(this, labelCol, featuresCol, weightCol, predictionCol, probabilityCol,
       rawPredictionCol, leafCol, impurity, numTrees, featureSubsetStrategy, maxDepth, maxBins,
       maxMemoryInMB, minInfoGain, pruneTree, minInstancesPerNode, minWeightFractionPerNode, seed,
-      subsamplingRate, thresholds, cacheNodeIds, checkpointInterval, bootstrap)
+      subsamplingRate, thresholds, cacheNodeIds, checkpointInterval, bootstrap,
+      intermediateStorageLevel)
 
     val trees = RandomForest
-      .run(instances, strategy, getNumTrees, getFeatureSubsetStrategy, getSeed, Some(instr))
+      .run(instances, strategy, getNumTrees, getFeatureSubsetStrategy, getSeed, Some(instr),
+        storageLevel = StorageLevel.fromString($(intermediateStorageLevel)))
       .map(_.asInstanceOf[DecisionTreeClassificationModel])
     trees.foreach(copyValues(_))
 

mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala

@@ -35,6 +35,7 @@ import org.apache.spark.mllib.tree.model.{GradientBoostedTreesModel => OldGBTMod
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types.StructType
+import org.apache.spark.storage.StorageLevel
 
 /**
  * <a href="http://en.wikipedia.org/wiki/Gradient_boosting">Gradient-Boosted Trees (GBTs)</a>
@@ -164,6 +165,10 @@ class GBTRegressor @Since("1.4.0") (@Since("1.4.0") override val uid: String)
   @Since("3.0.0")
   def setWeightCol(value: String): this.type = set(weightCol, value)
 
+  /** @group expertSetParam */
+  @Since("5.0.0")
+  def setIntermediateStorageLevel(value: String): this.type = set(intermediateStorageLevel, value)
+
   override protected def train(dataset: Dataset[_]): GBTRegressionModel = instrumented { instr =>
     val withValidation = isDefined(validationIndicatorCol) && $(validationIndicatorCol).nonEmpty
     val (trainDataset, validationDataset) = if (withValidation) {
@@ -178,16 +183,17 @@ class GBTRegressor @Since("1.4.0") (@Since("1.4.0") override val uid: String)
     instr.logParams(this, labelCol, featuresCol, predictionCol, leafCol, weightCol, impurity,
       lossType, maxDepth, maxBins, maxIter, maxMemoryInMB, minInfoGain, minInstancesPerNode,
       minWeightFractionPerNode, seed, stepSize, subsamplingRate, cacheNodeIds, checkpointInterval,
-      featureSubsetStrategy, validationIndicatorCol, validationTol)
+      featureSubsetStrategy, validationIndicatorCol, validationTol, intermediateStorageLevel)
 
     val categoricalFeatures = MetadataUtils.getCategoricalFeatures(dataset.schema($(featuresCol)))
     val boostingStrategy = super.getOldBoostingStrategy(categoricalFeatures, OldAlgo.Regression)
+    val storageLevel = StorageLevel.fromString($(intermediateStorageLevel))
     val (baseLearners, learnerWeights) = if (withValidation) {
       GradientBoostedTrees.runWithValidation(trainDataset, validationDataset, boostingStrategy,
-        $(seed), $(featureSubsetStrategy), Some(instr))
+        $(seed), $(featureSubsetStrategy), Some(instr), storageLevel)
     } else {
       GradientBoostedTrees.run(trainDataset, boostingStrategy,
-        $(seed), $(featureSubsetStrategy), Some(instr))
+        $(seed), $(featureSubsetStrategy), Some(instr), storageLevel)
     }
     baseLearners.foreach(copyValues(_))
 

mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala

@@ -36,6 +36,7 @@ import org.apache.spark.mllib.tree.model.{RandomForestModel => OldRandomForestMo
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions.{col, udf}
 import org.apache.spark.sql.types.StructType
+import org.apache.spark.storage.StorageLevel
 
 /**
  * <a href="http://en.wikipedia.org/wiki/Random_forest">Random Forest</a>
@@ -133,6 +134,10 @@ class RandomForestRegressor @Since("1.4.0") (@Since("1.4.0") override val uid: S
   @Since("3.0.0")
   def setWeightCol(value: String): this.type = set(weightCol, value)
 
+  /** @group expertSetParam */
+  @Since("5.0.0")
+  def setIntermediateStorageLevel(value: String): this.type = set(intermediateStorageLevel, value)
+
   override protected def train(
       dataset: Dataset[_]): RandomForestRegressionModel = instrumented { instr =>
     val categoricalFeatures: Map[Int, Int] =
@@ -154,10 +159,11 @@ class RandomForestRegressor @Since("1.4.0") (@Since("1.4.0") override val uid: S
     instr.logParams(this, labelCol, featuresCol, weightCol, predictionCol, leafCol, impurity,
       numTrees, featureSubsetStrategy, maxDepth, maxBins, maxMemoryInMB, minInfoGain,
       minInstancesPerNode, minWeightFractionPerNode, seed, subsamplingRate, cacheNodeIds,
-      checkpointInterval, bootstrap)
+      checkpointInterval, bootstrap, intermediateStorageLevel)
 
     val trees = RandomForest
-      .run(instances, strategy, getNumTrees, getFeatureSubsetStrategy, getSeed, Some(instr))
+      .run(instances, strategy, getNumTrees, getFeatureSubsetStrategy, getSeed, Some(instr),
+        storageLevel = StorageLevel.fromString($(intermediateStorageLevel)))
       .map(_.asInstanceOf[DecisionTreeRegressionModel])
     trees.foreach(copyValues(_))
 

mllib/src/main/scala/org/apache/spark/ml/tree/impl/GradientBoostedTrees.scala

@@ -49,18 +49,19 @@ private[spark] object GradientBoostedTrees extends Logging {
       boostingStrategy: OldBoostingStrategy,
       seed: Long,
       featureSubsetStrategy: String,
-      instr: Option[Instrumentation] = None):
+      instr: Option[Instrumentation] = None,
+      storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK):
         (Array[DecisionTreeRegressionModel], Array[Double]) = {
     val algo = boostingStrategy.treeStrategy.algo
     algo match {
       case OldAlgo.Regression =>
         GradientBoostedTrees.boost(input, input, boostingStrategy, validate = false,
-          seed, featureSubsetStrategy, instr)
+          seed, featureSubsetStrategy, instr, storageLevel)
       case OldAlgo.Classification =>
         // Map labels to -1, +1 so binary classification can be treated as regression.
         val remappedInput = input.map(x => Instance((x.label * 2) - 1, x.weight, x.features))
         GradientBoostedTrees.boost(remappedInput, remappedInput, boostingStrategy, validate = false,
-          seed, featureSubsetStrategy, instr)
+          seed, featureSubsetStrategy, instr, storageLevel)
       case _ =>
         throw new IllegalArgumentException(s"$algo is not supported by gradient boosting.")
     }
@@ -84,21 +85,22 @@ private[spark] object GradientBoostedTrees extends Logging {
       boostingStrategy: OldBoostingStrategy,
       seed: Long,
       featureSubsetStrategy: String,
-      instr: Option[Instrumentation] = None):
+      instr: Option[Instrumentation] = None,
+      storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK):
         (Array[DecisionTreeRegressionModel], Array[Double]) = {
     val algo = boostingStrategy.treeStrategy.algo
     algo match {
       case OldAlgo.Regression =>
         GradientBoostedTrees.boost(input, validationInput, boostingStrategy,
-          validate = true, seed, featureSubsetStrategy, instr)
+          validate = true, seed, featureSubsetStrategy, instr, storageLevel)
       case OldAlgo.Classification =>
         // Map labels to -1, +1 so binary classification can be treated as regression.
         val remappedInput = input.map(
           x => Instance((x.label * 2) - 1, x.weight, x.features))
         val remappedValidationInput = validationInput.map(
           x => Instance((x.label * 2) - 1, x.weight, x.features))
         GradientBoostedTrees.boost(remappedInput, remappedValidationInput, boostingStrategy,
-          validate = true, seed, featureSubsetStrategy, instr)
+          validate = true, seed, featureSubsetStrategy, instr, storageLevel)
       case _ =>
         throw new IllegalArgumentException(s"$algo is not supported by the gradient boosting.")
     }
@@ -294,7 +296,8 @@ private[spark] object GradientBoostedTrees extends Logging {
       validate: Boolean,
       seed: Long,
       featureSubsetStrategy: String,
-      instr: Option[Instrumentation] = None):
+      instr: Option[Instrumentation] = None,
+      storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK):
         (Array[DecisionTreeRegressionModel], Array[Double]) = {
     val earlyStopModelSizeThresholdInBytes = TreeConfig.trainingEarlyStopModelSizeThresholdInBytes
     lastEarlyStoppedModelSize = 0
@@ -324,7 +327,7 @@ private[spark] object GradientBoostedTrees extends Logging {
     // Prepare periodic checkpointers
     // Note: this is checkpointing the unweighted training error
     val predErrorCheckpointer = new PeriodicRDDCheckpointer[(Double, Double)](
-      treeStrategy.getCheckpointInterval(), sc, StorageLevel.MEMORY_AND_DISK)
+      treeStrategy.getCheckpointInterval(), sc, storageLevel)
 
     timer.stop("init")
 
@@ -349,7 +352,7 @@ private[spark] object GradientBoostedTrees extends Logging {
     // Cache input RDD for speedup during multiple passes.
     val treePoints = TreePoint.convertToTreeRDD(
       retaggedInput, splits, metadata)
-      .persist(StorageLevel.MEMORY_AND_DISK)
+      .persist(storageLevel)
       .setName("binned tree points")
 
     val firstCounts = BaggedPoint
@@ -359,7 +362,7 @@ private[spark] object GradientBoostedTrees extends Logging {
         require(bagged.subsampleCounts.length == 1)
         require(bagged.sampleWeight == bagged.datum.weight)
         bagged.subsampleCounts.head
-      }.persist(StorageLevel.MEMORY_AND_DISK)
+      }.persist(storageLevel)
       .setName("firstCounts at iter=0")
 
     val firstBagged = treePoints.zip(firstCounts)
@@ -372,7 +375,8 @@ private[spark] object GradientBoostedTrees extends Logging {
       metadata = metadata, bcSplits = bcSplits, strategy = treeStrategy, numTrees = 1,
       featureSubsetStrategy = featureSubsetStrategy, seed = seed, instr = instr,
       parentUID = None,
-      earlyStopModelSizeThresholdInBytes = earlyStopModelSizeThresholdInBytes)
+      earlyStopModelSizeThresholdInBytes = earlyStopModelSizeThresholdInBytes,
+      storageLevel = storageLevel)
       .head.asInstanceOf[DecisionTreeRegressionModel]
 
     firstCounts.unpersist()
@@ -397,11 +401,11 @@ private[spark] object GradientBoostedTrees extends Logging {
       timer.start("init validation")
       validationTreePoints = TreePoint.convertToTreeRDD(
         validationInput.retag(classOf[Instance]), splits, metadata)
-        .persist(StorageLevel.MEMORY_AND_DISK)
+        .persist(storageLevel)
       validatePredError = computeInitialPredictionAndError(
         validationTreePoints, firstTreeWeight, firstTreeModel, loss, bcSplits)
       validatePredErrorCheckpointer = new PeriodicRDDCheckpointer[(Double, Double)](
-        treeStrategy.getCheckpointInterval(), sc, StorageLevel.MEMORY_AND_DISK)
+        treeStrategy.getCheckpointInterval(), sc, storageLevel)
       validatePredErrorCheckpointer.update(validatePredError)
       bestValidateError = computeWeightedError(validationTreePoints, validatePredError)
       timer.stop("init validation")
@@ -437,7 +441,7 @@ private[spark] object GradientBoostedTrees extends Logging {
           // Update labels with pseudo-residuals
           val newLabel = -loss.gradient(pred, bagged.datum.label)
           (newLabel, bagged.subsampleCounts.head)
-        }.persist(StorageLevel.MEMORY_AND_DISK)
+        }.persist(storageLevel)
         .setName(s"labelWithCounts at iter=$m")
 
       val bagged = treePoints.zip(labelWithCounts)
@@ -451,7 +455,8 @@ private[spark] object GradientBoostedTrees extends Logging {
         metadata = metadata, bcSplits = bcSplits, strategy = treeStrategy,
         numTrees = 1, featureSubsetStrategy = featureSubsetStrategy,
         seed = seed + m, instr = None, parentUID = None,
-        earlyStopModelSizeThresholdInBytes = earlyStopModelSizeThresholdInBytes - accTreeSize)
+        earlyStopModelSizeThresholdInBytes = earlyStopModelSizeThresholdInBytes - accTreeSize,
+        storageLevel = storageLevel)
         .head.asInstanceOf[DecisionTreeRegressionModel]
 
       labelWithCounts.unpersist()

mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala

@@ -125,7 +125,8 @@ private[spark] object RandomForest extends Logging with Serializable {
       seed: Long,
       instr: Option[Instrumentation],
       parentUID: Option[String] = None,
-      earlyStopModelSizeThresholdInBytes: Long = 0): Array[DecisionTreeModel] = {
+      earlyStopModelSizeThresholdInBytes: Long = 0,
+      storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK): Array[DecisionTreeModel] = {
     lastEarlyStoppedModelSize = 0
     val timer = new TimeTracker()
     timer.start("total")
@@ -173,7 +174,7 @@ private[spark] object RandomForest extends Logging with Serializable {
       // At first, all the rows belong to the root nodes (node Id == 1).
       nodeIds = baggedInput.map { _ => Array.fill(numTrees)(1) }
       nodeIdCheckpointer = new PeriodicRDDCheckpointer[Array[Int]](
-        strategy.getCheckpointInterval(), sc, StorageLevel.MEMORY_AND_DISK)
+        strategy.getCheckpointInterval(), sc, storageLevel)
       nodeIdCheckpointer.update(nodeIds)
     }
 
@@ -307,7 +308,8 @@ private[spark] object RandomForest extends Logging with Serializable {
       featureSubsetStrategy: String,
       seed: Long,
       instr: Option[Instrumentation],
-      parentUID: Option[String] = None): Array[DecisionTreeModel] = {
+      parentUID: Option[String] = None,
+      storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK): Array[DecisionTreeModel] = {
     val earlyStopModelSizeThresholdInBytes = TreeConfig.trainingEarlyStopModelSizeThresholdInBytes
     val timer = new TimeTracker()
 
@@ -344,7 +346,7 @@ private[spark] object RandomForest extends Logging with Serializable {
         strategy.bootstrap,
         (tp: TreePoint) => tp.weight,
         seed = seed)
-      .persist(StorageLevel.MEMORY_AND_DISK)
+      .persist(storageLevel)
       .setName("bagged tree points")
 
     val trees = runBagged(
@@ -357,7 +359,8 @@ private[spark] object RandomForest extends Logging with Serializable {
       seed = seed,
       instr = instr,
       parentUID = parentUID,
-      earlyStopModelSizeThresholdInBytes = earlyStopModelSizeThresholdInBytes)
+      earlyStopModelSizeThresholdInBytes = earlyStopModelSizeThresholdInBytes,
+      storageLevel = storageLevel)
 
     baggedInput.unpersist()
     bcSplits.destroy()

mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala

@@ -340,7 +340,7 @@ private[spark] object TreeEnsembleParams {
  *
  * Note: Marked as private since this may be made public in the future.
  */
-private[ml] trait TreeEnsembleParams extends DecisionTreeParams {
+private[ml] trait TreeEnsembleParams extends DecisionTreeParams with HasIntermediateStorageLevel {
 
   /**
    * Fraction of the training data used for learning each decision tree, in range (0, 1].

mllib/src/test/scala/org/apache/spark/ml/classification/GBTClassifierSuite.scala

@@ -82,6 +82,19 @@ class GBTClassifierSuite extends MLTest with DefaultReadWriteTest {
     ParamsSuite.checkParams(model)
   }
 
+  test("SPARK-57870: intermediateStorageLevel param") {
+    val gbt = new GBTClassifier()
+    assert(gbt.getIntermediateStorageLevel === "MEMORY_AND_DISK")
+    gbt.setIntermediateStorageLevel("MEMORY_ONLY")
+    assert(gbt.getIntermediateStorageLevel === "MEMORY_ONLY")
+    intercept[IllegalArgumentException] {
+      new GBTClassifier().setIntermediateStorageLevel("NONE")
+    }
+    intercept[IllegalArgumentException] {
+      new GBTClassifier().setIntermediateStorageLevel("no_such_a_level")
+    }
+  }
+
   test("GBTClassifier: default params") {
     val gbt = new GBTClassifier
     assert(gbt.getLabelCol === "label")