Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[pyspark] Reuse the collective communicator. (#10198)

Browse Source
This commit is contained in:
Jiaming Yuan 2024-04-18 19:09:30 +08:00 committed by GitHub
parent 0aa2600399
commit 303c603c7d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
python-package/xgboost/spark/utils.py
View File

@ -14,7 +14,8 @@ import pyspark
from pyspark import BarrierTaskContext, SparkConf, SparkContext, SparkFiles, TaskContext
from pyspark.sql.session import SparkSession
from xgboost import Booster, XGBModel, collective
from xgboost import Booster, XGBModel
from xgboost.collective import CommunicatorContext as CCtx
from xgboost.tracker import RabitTracker
@ -42,22 +43,12 @@ def _get_default_params_from_func(
return filtered_params_dict
class CommunicatorContext:
"""A context controlling collective communicator initialization and finalization.
This isn't specificially necessary (note Part 3), but it is more understandable
coding-wise.
"""
class CommunicatorContext(CCtx): # pylint: disable=too-few-public-methods
"""Context with PySpark specific task ID."""
def __init__(self, context: BarrierTaskContext, **args: Any) -> None:
self.args = args
self.args["DMLC_TASK_ID"] = str(context.partitionId())
def __enter__(self) -> None:
collective.init(**self.args)
def __exit__(self, *args: Any) -> None:
collective.finalize()
args["DMLC_TASK_ID"] = str(context.partitionId())
super().__init__(**args)
def _start_tracker(context: BarrierTaskContext, n_workers: int) -> Dict[str, Any]: