maoli67660 opened a new pull request, #57206: URL: https://github.com/apache/spark/pull/57206
### What changes were proposed in this pull request? This is a sub-task of [SPARK-47103](https://issues.apache.org/jira/browse/SPARK-47103), which aims to make the storage level of MLlib's intermediate datasets configurable. It applies the shared param `HasIntermediateStorageLevel` (added in SPARK-57860, first reused in SPARK-57910) to the tree ensemble estimators: `RandomForestClassifier`, `RandomForestRegressor`, `GBTClassifier`, and `GBTRegressor`. - Scala: - `TreeEnsembleParams` now extends `HasIntermediateStorageLevel`, so all four estimators (and their models, which share the params traits) expose the param; each estimator gains a `setIntermediateStorageLevel` setter and logs the param via `Instrumentation`. - The storage level is threaded through the internal implementations with a `storageLevel` parameter defaulting to `StorageLevel.MEMORY_AND_DISK` (so the old `spark.mllib` API and `DecisionTreeClassifier`/`DecisionTreeRegressor`, which reach the same code, are unchanged): - `RandomForest.run` / `RandomForest.runBagged`: the `bagged tree points` RDD and the node-id-cache `PeriodicRDDCheckpointer`. - `GradientBoostedTrees.run` / `runWithValidation` / `boost`: the `binned tree points`, `firstCounts`, `labelWithCounts`, and validation `TreePoint` RDDs, plus the two prediction-error `PeriodicRDDCheckpointer`s. - Python: `_TreeEnsembleParams` mixes in `HasIntermediateStorageLevel`, and the four estimators gain `setIntermediateStorageLevel`, mirroring the Scala hierarchy so Scala/Python param parity is preserved for both estimators and models. ### Why are the changes needed? The tree ensemble trainers persist several intermediate RDDs internally with a hardcoded `MEMORY_AND_DISK` level. These datasets are created inside the algorithm, so users have no way to change their storage level, unlike the input `DataFrame` which they can cache themselves. Making this configurable (e.g. `DISK_ONLY`) improves resilience to executor loss: since SPARK-27677 the External Shuffle Service can serve disk-persisted cached blocks. `ALS` (SPARK-57910) and `KMeans` (SPARK-57860) already expose the same param; this PR extends it to the tree ensembles. ### Does this PR introduce _any_ user-facing change? Yes. `RandomForestClassifier`, `RandomForestRegressor`, `GBTClassifier`, and `GBTRegressor` (Scala and PySpark) gain a new expert param `intermediateStorageLevel` and a `setIntermediateStorageLevel` setter. The default is `"MEMORY_AND_DISK"`, so **behavior is unchanged unless the user sets it**. Before (no way to change the intermediate storage level): ```python rf = RandomForestClassifier(numTrees=100) # intermediate data always MEMORY_AND_DISK ``` After: ```python rf = RandomForestClassifier(numTrees=100).setIntermediateStorageLevel("DISK_ONLY") ``` ### How was this patch tested? - Each of the four estimator suites asserts the param's default value, that it can be set, and that invalid values (`"NONE"` and non-existent levels) are rejected. - `RandomForestClassifierSuite` and `GBTRegressorSuite` add an end-to-end test (modeled after `ALSStorageSuite`) that fits with `intermediateStorageLevel = "DISK_ONLY"` and verifies via a `SparkListener` that the intermediate RDDs were actually persisted at that level, covering both the `RandomForest` and `GradientBoostedTrees` code paths. - All four estimator suites pass (70/70), and the impl suites `ml.tree.impl.RandomForestSuite` / `GradientBoostedTreesSuite` pass (22/22) locally. - PySpark param parity is covered by the existing `pyspark.ml.tests.test_param` (`test_java_params`), which passes locally. - `dev/mima` reports no binary compatibility problems; no `MimaExcludes` entries are needed (the param members are new, unlike the ALS case in SPARK-57910). ### Was this patch authored or co-authored using generative AI tooling? Cooperate with: Claude Code (Fable 5) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
