text: # Copyright 2013-2022 Lawrence Livermore National Security, LLC and other
  # Spack Project Developers. See the top-level COPYRIGHT file for details.
  #
  # SPDX-License-Identifier: (Apache-2.0 OR MIT)
  
  
  from spack.package import *
  
  
  class PyHorovod(PythonPackage, CudaPackage):
      """Horovod is a distributed deep learning training framework for
      TensorFlow, Keras, PyTorch, and Apache MXNet."""
  
      homepage = "https://github.com/horovod"
      git      = "https://github.com/horovod/horovod.git"
  
      maintainers = ['adamjstewart', 'aweits', 'tgaddair']
  
      version('master', branch='master', submodules=True)
      version('0.21.3', tag='v0.21.3', submodules=True)
      version('0.21.2', tag='v0.21.2', submodules=True)
      version('0.21.1', tag='v0.21.1', submodules=True)
      version('0.21.0', tag='v0.21.0', submodules=True)
      version('0.20.3', tag='v0.20.3', submodules=True)
      version('0.20.2', tag='v0.20.2', submodules=True)
      version('0.20.1', tag='v0.20.1', submodules=True)
      version('0.20.0', tag='v0.20.0', submodules=True)
      version('0.19.5', tag='v0.19.5', submodules=True)
      version('0.19.4', tag='v0.19.4', submodules=True)
      version('0.19.3', tag='v0.19.3', submodules=True)
      version('0.19.2', tag='v0.19.2', submodules=True)
      version('0.19.1', tag='v0.19.1', submodules=True)
      version('0.19.0', tag='v0.19.0', submodules=True)
      version('0.18.2', tag='v0.18.2', submodules=True)
      version('0.18.1', tag='v0.18.1', submodules=True)
      version('0.18.0', tag='v0.18.0', submodules=True)
      version('0.17.1', tag='v0.17.1', submodules=True)
      version('0.17.0', tag='v0.17.0', submodules=True)
      version('0.16.4', tag='v0.16.4', submodules=True)
      version('0.16.3', tag='v0.16.3', submodules=True)
      version('0.16.2', tag='v0.16.2', submodules=True)
  
      # https://github.com/horovod/horovod/blob/master/docs/install.rst
      variant('frameworks', default='pytorch',
              description='Deep learning frameworks to build support for',
              values=('tensorflow', 'pytorch', 'mxnet', 'keras', 'spark', 'ray'),
              multi=True)
      variant('controllers', default='mpi',
              description='Controllers to coordinate work between processes',
              values=('mpi', 'gloo'), multi=True)
      variant('tensor_ops', default='nccl',
              description='Framework to use for GPU/CPU operations',
              values=('nccl', 'mpi', 'gloo', 'ccl'), multi=False)
      variant('cuda', default=True, description='Build with CUDA')
      variant('rocm', default=False, description='Build with ROCm')
  
      # Required dependencies
      depends_on('python@3.6:',    type=('build', 'run'), when='@0.20:')
      depends_on('py-setuptools',  type='build')
      depends_on('py-cloudpickle', type=('build', 'run'))
      depends_on('py-psutil',      type=('build', 'run'))
      depends_on('py-pyyaml',      type=('build', 'run'))
      depends_on('py-six',         type=('build', 'run'), when='@:0.19')
      depends_on('py-dataclasses', type=('build', 'run'), when='@0.20: ^python@:3.6')
  
      # Framework dependencies
      depends_on('py-tensorflow@1.1.0:',  type=('build', 'link', 'run'), when='frameworks=tensorflow')
      depends_on('py-tensorflow@1.15:',   type=('build', 'link', 'run'), when='frameworks=tensorflow @0.20:')
      depends_on('py-tensorflow-estimator', type=('build', 'run'), when='frameworks=tensorflow')
      depends_on('py-torch@0.4.0:',       type=('build', 'link', 'run'), when='frameworks=pytorch')
      depends_on('py-torch@1.2:',         type=('build', 'link', 'run'), when='frameworks=pytorch @0.20:')
      depends_on('py-torchvision',        type=('build', 'run'),         when='frameworks=pytorch @:0.19.1')
      depends_on('py-cffi@1.4.0:',        type=('build', 'run'),         when='frameworks=pytorch')
      depends_on('mxnet@1.4.1:+python',   type=('build', 'link', 'run'), when='frameworks=mxnet')
      depends_on('py-keras@2.0.8,2.1.2:', type=('build', 'run'),         when='frameworks=keras')
      depends_on('py-h5py@:2',        type=('build', 'run'),         when='frameworks=spark')
      depends_on('py-numpy',              type=('build', 'run'),         when='frameworks=spark')
      depends_on('py-petastorm@0.8.2',    type=('build', 'run'),         when='frameworks=spark @:0.19.1')
      depends_on('py-petastorm@0.9.0:',   type=('build', 'run'),         when='frameworks=spark @0.19.2:0.21.0')
      depends_on('py-petastorm@0.9.8:',   type=('build', 'run'),         when='frameworks=spark @0.21.1:')
      depends_on('py-pyarrow@0.15.0:',    type=('build', 'run'),         when='frameworks=spark')
      depends_on('py-pyspark@2.3.2:',     type=('build', 'run'),         when='frameworks=spark ^python@:3.7')
      depends_on('py-pyspark@3.0.0:',     type=('build', 'run'),         when='frameworks=spark ^python@3.8:')
      depends_on('py-ray',                type=('build', 'run'),         when='frameworks=ray')
  
      # Build dependencies
      depends_on('cmake@2.8.12:', type='build', when='@0.20:')
      depends_on('pkgconfig', type='build')
  
      # Controller dependencies
      depends_on('mpi', when='controllers=mpi')
      # There does not appear to be a way to use an external Gloo installation
      depends_on('cmake', type='build', when='controllers=gloo')
      depends_on('libuv@1.26:', when='controllers=gloo platform=darwin')
  
      # Tensor Operations dependencies
      depends_on('nccl@2:', when='tensor_ops=nccl')
      depends_on('mpi', when='tensor_ops=mpi')
      # There does not appear to be a way to use an external Gloo installation
      depends_on('cmake', type='build', when='tensor_ops=gloo')
  
      conflicts('cuda_arch=none', when='+cuda',
                msg='Must specify CUDA compute capabilities of your GPU, see '
                'https://developer.nvidia.com/cuda-gpus')
      conflicts('tensor_ops=nccl', when='~cuda~rocm', msg='NCCL requires either CUDA or ROCm support')
      conflicts('frameworks=ray', when='@:0.19', msg='Ray integration was added in 0.20.X')
      conflicts('controllers=gloo', when='@:0.20.0 platform=darwin', msg='Gloo cannot be compiled on MacOS')
  
      # https://github.com/horovod/horovod/pull/1835
      patch('fma.patch', when='@0.19.0:0.19.1')
  
      @property
      def import_modules(self):
          modules = [
              'horovod', 'horovod.runner', 'horovod.runner.util',
              'horovod.runner.elastic', 'horovod.runner.driver',
              'horovod.runner.common', 'horovod.runner.common.util',
              'horovod.runner.common.service', 'horovod.runner.http',
              'horovod.runner.task', 'horovod.common'
          ]
  
          if 'frameworks=tensorflow' in self.spec:
              modules.append('horovod.tensorflow')
  
          if 'frameworks=pytorch' in self.spec:
              modules.extend([
                  'horovod.torch', 'horovod.torch.mpi_lib',
                  'horovod.torch.elastic', 'horovod.torch.mpi_lib_impl'
              ])
  
          if 'frameworks=mxnet' in self.spec:
              modules.append('horovod.mxnet')
  
          if 'frameworks=keras' in self.spec:
              modules.extend(['horovod.keras', 'horovod._keras'])
  
          if 'frameworks=spark' in self.spec:
              modules.extend([
                  'horovod.spark', 'horovod.spark.driver',
                  'horovod.spark.common', 'horovod.spark.task'
              ])
  
          if 'frameworks=ray' in self.spec:
              modules.append('horovod.ray')
  
          if 'frameworks=tensorflow,keras' in self.spec:
              modules.append('horovod.tensorflow.keras')
  
          if 'frameworks=spark,pytorch' in self.spec:
              modules.append('horovod.spark.torch')
  
          if 'frameworks=spark,keras' in self.spec:
              modules.append('horovod.spark.keras')
  
          return modules
  
      def setup_build_environment(self, env):
          # https://github.com/horovod/horovod/blob/master/docs/install.rst#environment-variables
  
          # Build system
          env.set('PKG_CONFIG_EXECUTABLE',
                  self.spec['pkgconfig'].prefix.bin.join('pkg-config'))
          if '^cmake' in self.spec:
              env.set('HOROVOD_CMAKE', self.spec['cmake'].command.path)
          env.set('MAKEFLAGS', '-j{0}'.format(make_jobs))
  
          # Frameworks
          if 'frameworks=tensorflow' in self.spec:
              env.set('HOROVOD_WITH_TENSORFLOW', 1)
          else:
              env.set('HOROVOD_WITHOUT_TENSORFLOW', 1)
          if 'frameworks=pytorch' in self.spec:
              env.set('HOROVOD_WITH_PYTORCH', 1)
          else:
              env.set('HOROVOD_WITHOUT_PYTORCH', 1)
          if 'frameworks=mxnet' in self.spec:
              env.set('HOROVOD_WITH_MXNET', 1)
              env.set('MXNET_INCLUDE_PATH', self.spec['mxnet'].prefix.include)
              env.set('MXNET_LIBRARY_PATH', join_path(self.spec['mxnet'].libs[0]))
          else:
              env.set('HOROVOD_WITHOUT_MXNET', 1)
  
          # Controllers
          if 'controllers=mpi' in self.spec:
              env.set('HOROVOD_WITH_MPI', 1)
          else:
              env.set('HOROVOD_WITHOUT_MPI', 1)
          if 'controllers=gloo' in self.spec:
              env.set('HOROVOD_WITH_GLOO', 1)
          else:
              env.set('HOROVOD_WITHOUT_GLOO', 1)
  
          # Tensor Operations
          if 'tensor_ops=nccl' in self.spec:
              env.set('HOROVOD_GPU_ALLREDUCE', 'NCCL')
              env.set('HOROVOD_GPU_ALLGATHER', 'NCCL')
              env.set('HOROVOD_GPU_BROADCAST', 'NCCL')
  
              env.set('HOROVOD_NCCL_HOME', self.spec['nccl'].prefix)
              env.set('HOROVOD_NCCL_INCLUDE',
                      self.spec['nccl'].headers.directories[0])
              env.set('HOROVOD_NCCL_LIB', self.spec['nccl'].libs.directories[0])
  
              if '+cuda' in self.spec:
                  env.set('HOROVOD_GPU', 'CUDA')
  
                  env.set('HOROVOD_CUDA_HOME', self.spec['cuda'].prefix)
                  cuda_cc_list = ','.join(self.spec.variants['cuda_arch'].value)
                  env.set('HOROVOD_BUILD_CUDA_CC_LIST', cuda_cc_list)
                  env.set('HOROVOD_CUDA_INCLUDE',
                          self.spec['cuda'].headers.directories[0])
                  env.set('HOROVOD_CUDA_LIB',
                          self.spec['cuda'].libs.directories[0])
              elif '+rocm' in self.spec:
                  env.set('HOROVOD_GPU', 'ROCM')
                  # env.set('HOROVOD_ROCM_HOME', self.spec['rocm'].prefix)
          else:
              env.set('HOROVOD_CPU_OPERATIONS',
                      self.spec.variants['tensor_ops'].value.upper())
  
      def test(self):
          super(PyHorovod, self).test()
          run_test(self.prefix.bin.horovodrun, '--check-build')
  