rank = _basics.rank local_rank = _basics.local_rank cross_rank = _basics.cross_rank mpi_threads_supported = _basics.mpi_threads_supported mpi_enabled = _basics.mpi_enabled mpi_built = _basics.mpi_built gloo_enabled = _basics.gloo_enabled gloo_built = _basics.gloo_built nccl_built = _basics.nccl_built ddl_built = _basics.ddl_built ccl_built = _basics.ccl_built cuda_built = _basics.cuda_built rocm_built = _basics.rocm_built dll_path = os.path.join(os.path.dirname(__file__), 'mpi_lib' + get_ext_suffix()) MPI_MXNET_LIB_CTYPES = ctypes.CDLL(dll_path, ctypes.RTLD_GLOBAL) def allreduce(tensor, average=True, name=None, priority=0, prescale_factor=1.0, postscale_factor=1.0): """ A function that performs averaging or summation of the input tensor over all the Horovod processes. The input tensor is not modified. The reduction operation is keyed by the name. If name is not provided, an incremented auto-generated name is used. The tensor type and shape must be
def _load_library(name): """Loads a .so file containing the specified operators. Args: name: The name of the .so file to load. Raises: NotFoundError if were not able to load .so file. """ filename = resource_loader.get_path_to_datafile(name) library = load_library.load_op_library(filename) return library MPI_LIB = _load_library('mpi_lib' + get_ext_suffix()) _basics = _HorovodBasics(__file__, 'mpi_lib') # import basic methods init = _basics.init shutdown = _basics.shutdown is_initialized = _basics.is_initialized start_timeline = _basics.start_timeline stop_timeline = _basics.stop_timeline size = _basics.size local_size = _basics.local_size rank = _basics.rank local_rank = _basics.local_rank mpi_threads_supported = _basics.mpi_threads_supported mpi_enabled = _basics.mpi_enabled
"""Loads a .so file containing the specified operators. Args: name: The name of the .so file to load. Raises: NotFoundError if were not able to load .so file. """ filename = resource_loader.get_path_to_datafile(name) library = load_library.load_op_library(filename) return library # Check possible symbol not found error from tensorflow version mismatch try: print("get_ext_suffix() = {}".format(str(get_ext_suffix()))) MPI_LIB = _load_library('mpi_lib' + get_ext_suffix()) except Exception as e: check_installed_version('tensorflow', tf.__version__, e) raise e else: check_installed_version('tensorflow', tf.__version__) _basics = _HorovodBasics(__file__, 'mpi_lib') # import basic methods init = _basics.init shutdown = _basics.shutdown is_initialized = _basics.is_initialized start_timeline = _basics.start_timeline stop_timeline = _basics.stop_timeline