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 7 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
35 changes: 23 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 @@ -164,6 +164,12 @@ class _SparkXGBParams(
+ "Note: The auto repartitioning judgement is not fully accurate, so it is recommended"
+ "to have force_repartition be True.",
)
repartition_random_shuffle = Param(
Params._dummy(),
"repartition_random_shuffle",
"A boolean variable. Set repartition_random_shuffle=true if you want to random shuffle "
"dataset when repartitioning is required. By default is True."
)
feature_names = Param(
Params._dummy(), "feature_names", "A list of str to specify feature names."
)
Expand Down Expand Up @@ -270,15 +276,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 @@ -470,6 +467,7 @@ def __init__(self):
num_workers=1,
use_gpu=False,
force_repartition=False,
repartition_random_shuffle=True,
feature_names=None,
feature_types=None,
arbitrary_params_dict={},
Expand Down Expand Up @@ -690,8 +688,21 @@ 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.
if self.getOrDefault(self.repartition_random_shuffle):
# In some cases, spark round-robin repartition might cause data skew
# use random shuffle can address it.
dataset = dataset.repartition(num_workers, rand(1))
else:
dataset = dataset.repartition(num_workers)

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