# IT4Innovations 2023
# JK

easyblock = 'PythonBundle'

name = 'Horovod'
version = '0.27.0'
local_tf_version = '2.10.1'
local_cuda_suffix = '-CUDA-%(cudaver)s'
versionsuffix = local_cuda_suffix + '-TensorFlow-%s' % local_tf_version

homepage = 'https://github.com/uber/horovod'
description = "Horovod is a distributed training framework for TensorFlow."

toolchain = {'name': 'foss', 'version': '2022a'}

builddependencies = [
    ('CMake', '3.23.1'),
    ('flatbuffers', '2.0.7'),
]
dependencies = [
    ('CUDA', '11.7.0', '', SYSTEM),
    ('UCX-CUDA', '1.12.1', local_cuda_suffix),
    ('NCCL', '2.12.12', local_cuda_suffix),
    ('Python', '3.10.4'),
    ('PyYAML', '6.0'),
    ('TensorFlow', local_tf_version, local_cuda_suffix),
]

use_pip = True
sanity_pip_check = True

preinstallopts = 'HOROVOD_WITH_MPI=1 HOROVOD_GPU_ALLREDUCE=NCCL HOROVOD_GPU_BROADCAST=NCCL '
preinstallopts += 'HOROVOD_WITH_TENSORFLOW=1 HOROVOD_WITHOUT_PYTORCH=1 HOROVOD_WITHOUT_MXNET=1 '

# Bug in CMake causes a race condition on horovod_cuda_kernels_generated_cuda_kernels.cu.o.NVCC-depend
maxparallel = 1

exts_list = [
    ('cloudpickle', '2.2.1', {
        'checksums': ['d89684b8de9e34a2a43b3460fbca07d09d6e25ce858df4d5a44240403b6178f5'],
    }),
    ('horovod', version, {
        'checksums': ['8a76a98266579de46a1b1cde53d5ec115fc163ffdaa65a8f771aee936e804b33'],
    }),
]

sanity_check_paths = {
    'files': ['bin/horovodrun'],
    'dirs': ['lib/python%(pyshortver)s/site-packages'],
}

moduleclass = 'tools'