Add Python binding for rabit ops. (#5743)
This commit is contained in:
@@ -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.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user