wehub-resource-sync
22 lines
711 B
Python
22 lines
711 B
Python
from contextlib import contextmanager
|
|
|
|
from ray.train.v2._internal.execution.train_fn_utils import get_train_fn_utils
|
|
from ray.train.xgboost.config import XGBoostConfig as XGBoostConfigV1
|
|
|
|
|
|
class XGBoostConfig(XGBoostConfigV1):
|
|
@property
|
|
def train_func_context(self):
|
|
distributed_context = super(XGBoostConfig, self).train_func_context
|
|
|
|
@contextmanager
|
|
def collective_communication_context():
|
|
# The distributed_context is only needed in distributed mode
|
|
if get_train_fn_utils().is_distributed():
|
|
with distributed_context():
|
|
yield
|
|
else:
|
|
yield
|
|
|
|
return collective_communication_context
|