# Copyright 2019 Uber Technologies, Inc. All Rights Reserved. # Modifications copyright Microsoft # Modifications copyright (C) 2020, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== import os import subprocess import sys import textwrap from setuptools import setup, Extension, find_packages from setuptools.command.build_ext import build_ext from horovod import __version__ class CMakeExtension(Extension): def __init__(self, name, cmake_lists_dir='.', sources=[], **kwa): Extension.__init__(self, name, sources=sources, **kwa) self.cmake_lists_dir = os.path.abspath(cmake_lists_dir) tensorflow_mpi_lib = CMakeExtension('horovod.tensorflow.mpi_lib', cmake_lists_dir='.', sources=[]) torch_mpi_lib_v2 = CMakeExtension('horovod.torch.mpi_lib_v2', cmake_lists_dir='.', sources=[]) mxnet_mpi_lib = CMakeExtension('horovod.mxnet.mpi_lib', cmake_lists_dir='.', sources=[]) def is_build_action(): if len(sys.argv) <= 1: return False if sys.argv[1].startswith('build'): return True if sys.argv[1].startswith('bdist'): return True if sys.argv[1].startswith('install'): return True def get_cmake_bin(): return os.environ.get('HOROVOD_CMAKE', 'cmake') class custom_build_ext(build_ext): def build_extensions(self): cmake_bin = get_cmake_bin() config = 'Debug' if self.debug else 'Release' ext_name = self.extensions[0].name build_dir = self.get_ext_fullpath(ext_name).replace(self.get_ext_filename(ext_name), '') build_dir = os.path.abspath(build_dir) cmake_args = ['-DCMAKE_BUILD_TYPE=' + config, '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}'.format(config.upper(), build_dir), '-DPYTHON_EXECUTABLE:FILEPATH=' + sys.executable] make_args = ['-j8'] if not os.environ.get('MAKEFLAGS') else [] if self.verbose: make_args.append('VERBOSE=1') cmake_build_args = ['--config', config] if make_args: # -- specifies that these args are going to the native build tool: make cmake_build_args += ['--'] + make_args if not os.path.exists(self.build_temp): os.makedirs(self.build_temp) # Config and build the extension try: subprocess.check_call([cmake_bin, self.extensions[0].cmake_lists_dir] + cmake_args, cwd=self.build_temp) subprocess.check_call([cmake_bin, '--build', '.'] + cmake_build_args, cwd=self.build_temp) except OSError as e: raise RuntimeError('CMake failed: {}'.format(str(e))) # python packages required to use horovod in general require_list = ['cloudpickle', 'psutil', 'pyyaml', 'dataclasses;python_version<"3.7"'] # python packages required / recommended to develop horovod # e.g., set of framework versions pinned for development, keep in sync with Dockerfile.test.cpu # NOTE: do not use versions with +cpu or +gpu here as users would need to add --find-links to pip dev_require_list = ['tensorflow-cpu==1.15.0', 'keras==2.2.4', 'torch==1.2.0', 'torchvision==0.4.0', 'mxnet==1.5.0', 'pyspark==2.4.0'] # python packages required only to run tests test_require_list = ['mock', 'pytest', 'pytest-forked', 'parameterized'] # framework dependencies tensorflow_require_list = ['tensorflow'] tensorflow_cpu_require_list = ['tensorflow-cpu'] tensorflow_gpu_require_list = ['tensorflow-gpu'] keras_require_list = ['keras>=2.0.8,!=2.0.9,!=2.1.0,!=2.1.1'] pytorch_require_list = ['torch'] mxnet_require_list = ['mxnet>=1.4.1'] pyspark_require_list = ['pyspark>=2.3.2;python_version<"3.8"', # TODO: change to 'pyspark>=3.0.0' once spark3 is released 'pyspark>=3.0.0.dev;python_version>="3.8"'] spark_require_list = ['h5py>=2.9', 'numpy', 'petastorm>=0.9.0,!=0.9.3', 'pyarrow>=0.15.0'] + \ pyspark_require_list ray_require_list = ['ray'] # all frameworks' dependencies all_frameworks_require_list = tensorflow_require_list + \ tensorflow_gpu_require_list + \ keras_require_list + \ pytorch_require_list + \ mxnet_require_list + \ spark_require_list # Skip cffi if pytorch extension explicitly disabled if not os.environ.get('HOROVOD_WITHOUT_PYTORCH'): require_list.append('cffi>=1.4.0') def get_package_version(): return __version__ + "+" + os.environ['HOROVOD_LOCAL_VERSION'] if 'HOROVOD_LOCAL_VERSION' in os.environ else __version__ setup(name='horovod', version=get_package_version(), packages=find_packages(), description='Distributed training framework for TensorFlow, Keras, PyTorch, and Apache MXNet.', author='The Horovod Authors', long_description=textwrap.dedent('''\ Horovod is a distributed training framework for TensorFlow, Keras, PyTorch, and Apache MXNet. The goal of Horovod is to make distributed Deep Learning fast and easy to use.'''), url='https://github.com/horovod/horovod', keywords=['deep learning', 'tensorflow', 'keras', 'pytorch', 'mxnet', 'spark', 'AI'], classifiers=[ 'License :: OSI Approved :: Apache Software License', 'Development Status :: 4 - Beta', 'Intended Audience :: Developers', 'Topic :: Scientific/Engineering :: Artificial Intelligence', ], ext_modules=[tensorflow_mpi_lib, torch_mpi_lib_v2, mxnet_mpi_lib], cmdclass={'build_ext': custom_build_ext}, # cffi is required for PyTorch # If cffi is specified in setup_requires, it will need libffi to be installed on the machine, # which is undesirable. Luckily, `install` action will install cffi before executing build, # so it's only necessary for `build*` or `bdist*` actions. setup_requires=require_list if is_build_action() else [], install_requires=require_list, tests_require=test_require_list, extras_require={ 'all-frameworks': all_frameworks_require_list, 'tensorflow': tensorflow_require_list, 'tensorflow-cpu': tensorflow_cpu_require_list, 'tensorflow-gpu': tensorflow_gpu_require_list, 'keras': keras_require_list, 'pytorch': pytorch_require_list, 'mxnet': mxnet_require_list, 'spark': spark_require_list, 'ray': ray_require_list, 'dev': dev_require_list, }, # not used by pip since 19.0: https://github.com/pypa/pip/issues/4187#issuecomment-415067034 # here for completeness as pip install needs some of these via -f for versions with '+cpu' # for examples, see Dockerfile.test.cpu and Dockerfile.test.gpu dependency_links=[ 'https://download.pytorch.org/whl/torch_stable.html', 'https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html', 'https://dist.mxnet.io/python/all' ], python_requires='>=3.6', zip_safe=False, entry_points={ 'console_scripts': [ 'horovodrun = horovod.runner.launch:run_commandline' ] })