Skip to content

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

Open
maoli67660 wants to merge 1 commit into
apache:masterfrom
maoli67660:SPARK-57870
Open

[SPARK-57870][ML][PYTHON] Apply HasIntermediateStorageLevel to tree ensemble estimators (GBT, RandomForest)#57206
maoli67660 wants to merge 1 commit into
apache:masterfrom
maoli67660:SPARK-57870

Conversation

@maoli67660

Copy link
Copy Markdown
Contributor

What changes were proposed in this pull request?

This is a sub-task of 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 PeriodicRDDCheckpointers.
  • 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):

rf = RandomForestClassifier(numTrees=100)   # intermediate data always MEMORY_AND_DISK

After:

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)

…nsemble estimators (GBT, RandomForest)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant