view release on metacpan or  search on metacpan

xgboost/python-package/xgboost/rabit.py  view on Meta::CPAN

# coding: utf-8
# pylint: disable= invalid-name

"""Distributed XGBoost Rabit related API."""
from __future__ import absolute_import
import sys
import ctypes
import numpy as np

from .core import _LIB, c_str, STRING_TYPES
from .compat import pickle


def _init_rabit():
    """internal library initializer."""
    if _LIB is not None:
        _LIB.RabitGetRank.restype = ctypes.c_int
        _LIB.RabitGetWorldSize.restype = ctypes.c_int
        _LIB.RabitIsDistributed.restype = ctypes.c_int
        _LIB.RabitVersionNumber.restype = ctypes.c_int


def init(args=None):
    """Initialize the rabit library with arguments"""
    if args is None:
        args = []
    arr = (ctypes.c_char_p * len(args))()
    arr[:] = args
    _LIB.RabitInit(len(arr), arr)


def finalize():
    """Finalize the process, notify tracker everything is done."""
    _LIB.RabitFinalize()


def get_rank():
    """Get rank of current process.

    Returns
    -------
    rank : int
        Rank of current process.
    """
    ret = _LIB.RabitGetRank()
    return ret


def get_world_size():
    """Get total number workers.

    Returns
    -------
    n : int
        Total number of process.
    """
    ret = _LIB.RabitGetWorldSize()
    return ret


def tracker_print(msg):