Add Python binding for rabit ops. (#5743)

This commit is contained in:
Jiaming Yuan
2020-06-02 19:47:23 +08:00
committed by GitHub
parent e533908922
commit e49607af19
2 changed files with 53 additions and 2 deletions

View File

@@ -56,6 +56,12 @@ def get_world_size():
return ret
def is_distributed():
'''If rabit is distributed.'''
is_dist = _LIB.RabitIsDistributed()
return is_dist
def tracker_print(msg):
"""Print message to the tracker.
@@ -143,6 +149,14 @@ DTYPE_ENUM__ = {
}
class Op: # pylint: disable=too-few-public-methods
'''Supported operations for rabit.'''
MAX = 0
MIN = 1
SUM = 2
OR = 3
def allreduce(data, op, prepare_fun=None):
"""Perform allreduce, return the result.