Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[pyspark] Fix xgboost spark estimator dataset repartition issues #8231

Merged
merged 8 commits into from Sep 22, 2022
Merged

[pyspark] Fix xgboost spark estimator dataset repartition issues #8231

Changes from 6 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
9854131
update
WeichenXu123 Sep 8, 2022
7bcdda3
add debug log
WeichenXu123 Sep 9, 2022
ec05b26
update
WeichenXu123 Sep 10, 2022
9f8f64e
Merge branch 'master' into fix-repartition
WeichenXu123 Sep 13, 2022
8596b89
remove debug logging
WeichenXu123 Sep 13, 2022
5f89620
remove debug log
WeichenXu123 Sep 14, 2022
9410a89
update
WeichenXu123 Sep 15, 2022
ba45f07
fix lint
Sep 20, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
26 changes: 14 additions & 12 deletions python-package/xgboost/spark/core.py
View file
Open in desktop
Expand Up @@ -20,7 +20,7 @@
HasWeightCol,
)
from pyspark.ml.util import MLReadable, MLWritable
from pyspark.sql.functions import col, countDistinct, pandas_udf, struct
from pyspark.sql.functions import col, countDistinct, pandas_udf, struct, rand
from pyspark.sql.types import (
ArrayType,
DoubleType,
Expand Down Expand Up @@ -270,15 +270,6 @@ def _validate_params(self):
f"It cannot be less than 1 [Default is 1]"
)

if (
self.getOrDefault(self.force_repartition)
and self.getOrDefault(self.num_workers) == 1
):
get_logger(self.__class__.__name__).warning(
"You set force_repartition to true when there is no need for a repartition."
"Therefore, that parameter will be ignored."
)

if self.getOrDefault(self.features_cols):
if not self.getOrDefault(self.use_gpu):
raise ValueError("features_cols param requires enabling use_gpu.")
Expand Down Expand Up @@ -690,8 +681,19 @@ def _fit(self, dataset):
num_workers,
)

if self._repartition_needed(dataset):
dataset = dataset.repartition(num_workers)
if self._repartition_needed(dataset) or (
self.isDefined(self.validationIndicatorCol)
and self.getOrDefault(self.validationIndicatorCol)
):
# If validationIndicatorCol defined, we always repartition dataset
# to balance data, because user might unionise train and validation dataset,
# without shuffling data then some partitions might contain only train or validation
# dataset.
# Repartition on `rand` column to avoid repartition
# result unbalance. Directly using `.repartition(N)` might result in some
WeichenXu123 marked this conversation as resolved.
Show resolved Hide resolved
# empty partitions.
dataset = dataset.repartition(num_workers, rand(1))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I still prefer to add a parameter to control if random shuffle. Since my PR #8245 can detect if the training data is empty DMatrix, so we can add some log printing about data skew and how to resolve it by enabling xxx parameter. Does that make sense?


train_params = self._get_distributed_train_params(dataset)
booster_params, train_call_kwargs_params = self._get_xgb_train_call_args(
train_params
Expand Down