[pyspark] Allow to avoid repartition #10408
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
trivialfis
requested changes
Jun 11, 2024
python-package/xgboost/spark/core.py
Outdated
Comment on lines
709
to
716
| assert dataset._sc._jvm is not None | ||
| query_plan = dataset._sc._jvm.PythonSQLUtils.explainString( | ||
| dataset._jdf.queryExecution(), "extended" | ||
| ) | ||
| start = query_plan.index("== Optimized Logical Plan ==") | ||
| start += len("== Optimized Logical Plan ==") + 1 | ||
|
|
||
| query_plan[start : start + len("Repartition")] == "Repartition" |
There was a problem hiding this comment.
Please remove the commented code.
doc/tutorials/spark_estimator.rst
Outdated
| XGBoost needs to repartition to the num_workers to ensure there will be num_workers training | ||
| tasks running at the same time, but repartition is a costly operation. To avoid repartition, | ||
| users can set ``spark.sql.files.maxPartitionNum`` and ``spark.sql.files.minPartitionNum`` | ||
| to num_workers. |
There was a problem hiding this comment.
Could you please elaborate on when a user might want to force reparittion dataset as well?
trivialfis
approved these changes
Jun 11, 2024
trivialfis
pushed a commit
to trivialfis/xgboost
that referenced
this pull request
Jun 12, 2024
trivialfis
pushed a commit
to trivialfis/xgboost
that referenced
this pull request
Jun 12, 2024
trivialfis
added a commit
that referenced
this pull request
Jun 13, 2024
Co-authored-by: Bobby Wang <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The current xgboost pyspark will not repartition the dataset only when the last operator of input dataset is repartition and the repartitioned number is equal to num_workers. Overall, this looks good. But we should also keep in mind that repartition is really expensive especially for the GPU case, the repartition probably dominates the most of training time.
There is a real xgboost case that reading from file directly using spark and then fit the data into xgboost estimator. So for this kind of case, we can make the partition number equal to num_workers by playing with some spark configurations, and the data partitions should have been well balanced by spark. So I think, it's safe for xgboost to skip the repartition internally which can really improve the whole xgboost end to end time.
And on the other hand, xgboost already supports the force_repartition, so if user really would like to enable the repartition, they can set force_repartition to true.