setup.py 8 KB
Newer Older
Jeff Rasley's avatar
Jeff Rasley committed
1
2
3
4
5
6
7
8
9
10
"""
Copyright 2020 The Microsoft DeepSpeed Team

DeepSpeed library

Create a new wheel via the following command: python setup.py bdist_wheel

The wheel will be located at: dist/*.whl
"""

11
import os
12
import shutil
13
14
import subprocess
import warnings
Jeff Rasley's avatar
Jeff Rasley committed
15
from setuptools import setup, find_packages
Samyam Rajbhandari's avatar
Samyam Rajbhandari committed
16
import time
17
18
19

try:
    import torch
401qingkong's avatar
401qingkong committed
20
21
    from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
    from torch.utils.hipify import hipify_python
22
23
24
except ImportError:
    raise ImportError('Unable to import torch, please visit https://pytorch.org/ '
                      'to see how to properly install torch on your system.')
401qingkong's avatar
401qingkong committed
25
26
27
28
29
###aiss add 
is_rocm_pytorch = False
if torch.__version__ >= '1.5':
    from torch.utils.cpp_extension import ROCM_HOME
    is_rocm_pytorch = True if ((torch.version.hip is not None) and (ROCM_HOME is not None)) else False
30

31
from op_builder import ALL_OPS, get_default_compute_capatabilities
32
33
34
35
36
37
38
39


def fetch_requirements(path):
    with open(path, 'r') as fd:
        return [r.strip() for r in fd.readlines()]


install_requires = fetch_requirements('requirements/requirements.txt')
40
41
42
43
44
extras_require = {
    '1bit_adam': fetch_requirements('requirements/requirements-1bit-adam.txt'),
    'readthedocs': fetch_requirements('requirements/requirements-readthedocs.txt'),
    'dev': fetch_requirements('requirements/requirements-dev.txt'),
}
401qingkong's avatar
401qingkong committed
45
46
47
48
49
###aiss add ################
if is_rocm_pytorch:
    print("NOTE: Please manually install torch and torchvision packages for ROCm")
    #install_requires = fetch_requirements('requirements/requirements-rocm.txt')

50

51
# If MPI is available add 1bit-adam requirements
401qingkong's avatar
401qingkong committed
52
53
54
55
56
57
##aiss add: cupy 9.0 has been setup manufully
#if torch.cuda.is_available():
#    if shutil.which('ompi_info') or shutil.which('mpiname'):
#        cupy = f"cupy-cuda{torch.version.cuda.replace('.','')[:3]}"
#        print("cupy version: ", cupy)
#        extras_require['1bit_adam'].append(cupy)
58
59
60
61
62
63
64

# Make an [all] extra that installs all needed dependencies
all_extras = set()
for extra in extras_require.items():
    for req in extra[1]:
        all_extras.add(req)
extras_require['all'] = list(all_extras)
Jeff Rasley's avatar
Jeff Rasley committed
65
66

cmdclass = {}
67
68

# For any pre-installed ops force disable ninja
Jeff Rasley's avatar
Jeff Rasley committed
69
cmdclass['build_ext'] = BuildExtension.with_options(use_ninja=False)
Jeff Rasley's avatar
Jeff Rasley committed
70

71
72
TORCH_MAJOR = torch.__version__.split('.')[0]
TORCH_MINOR = torch.__version__.split('.')[1]
73
74

if not torch.cuda.is_available():
75
    # Fix to allow docker builds, similar to https://github.com/NVIDIA/apex/issues/486
76
    print(
77
        "[WARNING] Torch did not find cuda available, if cross-compiling or running with cpu only "
78
79
80
        "you can ignore this message. Adding compute capability for Pascal, Volta, and Turing "
        "(compute capabilities 6.0, 6.1, 6.2)")
    if os.environ.get("TORCH_CUDA_ARCH_LIST", None) is None:
81
        os.environ["TORCH_CUDA_ARCH_LIST"] = get_default_compute_capatabilities()
401qingkong's avatar
401qingkong committed
82
83
84
85
86
87
88
###########aiss add ######################only need run once for convert
#if is_rocm_pytorch:
#    import shutil
#    this_dir = os.path.dirname(os.path.abspath(__file__))
#    hipify_python.hipify(project_directory=this_dir, output_directory=this_dir, includes="csrc/*",
#                                show_detailed=True, is_pytorch_extension=True)
#    print("cuda file has been transformed to hip format!!!")
89
90
ext_modules = []

91
# Default to pre-install kernels to false so we rely on JIT
401qingkong's avatar
401qingkong committed
92
93


94
95
BUILD_OP_DEFAULT = int(os.environ.get('DS_BUILD_OPS', 0))
print(f"DS_BUILD_OPS={BUILD_OP_DEFAULT}")
96
97
98


def command_exists(cmd):
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
    result = subprocess.Popen(f'type {cmd}', stdout=subprocess.PIPE, shell=True)
    return result.wait() == 0


def op_enabled(op_name):
    assert hasattr(ALL_OPS[op_name], 'BUILD_VAR'), \
        f"{op_name} is missing BUILD_VAR field"
    env_var = ALL_OPS[op_name].BUILD_VAR
    return int(os.environ.get(env_var, BUILD_OP_DEFAULT))


install_ops = dict.fromkeys(ALL_OPS.keys(), False)
for op_name, builder in ALL_OPS.items():
    op_compatible = builder.is_compatible()

    # If op is compatible update install reqs so it can potentially build/run later
    if op_compatible:
        reqs = builder.python_requirements()
        install_requires += builder.python_requirements()
401qingkong's avatar
401qingkong committed
118
119
 ######aiss debug###############
    print("op_enabled(op_name): , op_compatible: ", op_enabled(op_name), op_compatible)
120
121
122
    # If op install enabled, add builder to extensions
    if op_enabled(op_name) and op_compatible:
        install_ops[op_name] = op_enabled(op_name)
401qingkong's avatar
401qingkong committed
123
        ext_modules.append(builder.builder(is_rocm_pytorch))
124
125
126
127

compatible_ops = {op_name: op.is_compatible() for (op_name, op) in ALL_OPS.items()}

print(f'Install Ops={install_ops}')
128
129
130
131

# Write out version/git info
git_hash_cmd = "git rev-parse --short HEAD"
git_branch_cmd = "git rev-parse --abbrev-ref HEAD"
132
133
134
135
136
137
138
139
140
if command_exists('git') and 'DS_BUILD_STRING' not in os.environ:
    try:
        result = subprocess.check_output(git_hash_cmd, shell=True)
        git_hash = result.decode('utf-8').strip()
        result = subprocess.check_output(git_branch_cmd, shell=True)
        git_branch = result.decode('utf-8').strip()
    except subprocess.CalledProcessError:
        git_hash = "unknown"
        git_branch = "unknown"
141
142
143
else:
    git_hash = "unknown"
    git_branch = "unknown"
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

# Parse the DeepSpeed version string from version.txt
version_str = open('version.txt', 'r').read().strip()

# Build specifiers like .devX can be added at install time. Otherwise, add the git hash.
# example: DS_BUILD_STR=".dev20201022" python setup.py sdist bdist_wheel

# Building wheel for distribution, update version file
if 'DS_BUILD_STRING' in os.environ:
    # Build string env specified, probably building for distribution
    with open('build.txt', 'w') as fd:
        fd.write(os.environ.get('DS_BUILD_STRING'))
    version_str += os.environ.get('DS_BUILD_STRING')
elif os.path.isfile('build.txt'):
    # build.txt exists, probably installing from distribution
    with open('build.txt', 'r') as fd:
        version_str += fd.read().strip()
else:
    # None of the above, probably installing from source
    version_str += f'+{git_hash}'

torch_version = ".".join([TORCH_MAJOR, TORCH_MINOR])
166
167
168
169
# Set cuda_version to 0.0 if cpu-only
cuda_version = "0.0"
if torch.version.cuda is not None:
    cuda_version = ".".join(torch.version.cuda.split('.')[:2])
170
171
172
torch_info = {"version": torch_version, "cuda_version": cuda_version}

print(f"version={version_str}, git_hash={git_hash}, git_branch={git_branch}")
173
with open('deepspeed/git_version_info_installed.py', 'w') as fd:
174
    fd.write(f"version='{version_str}'\n")
175
176
    fd.write(f"git_hash='{git_hash}'\n")
    fd.write(f"git_branch='{git_branch}'\n")
Jeff Rasley's avatar
Jeff Rasley committed
177
    fd.write(f"installed_ops={install_ops}\n")
178
179
    fd.write(f"compatible_ops={compatible_ops}\n")
    fd.write(f"torch_info={torch_info}\n")
180
181

print(f'install_requires={install_requires}')
182
183
print(f'compatible_ops={compatible_ops}')
print(f'ext_modules={ext_modules}')
Jeff Rasley's avatar
Jeff Rasley committed
184

185
186
187
188
189
# Parse README.md to make long_description for PyPI page.
thisdir = os.path.abspath(os.path.dirname(__file__))
with open(os.path.join(thisdir, 'README.md'), encoding='utf-8') as fin:
    readme_text = fin.read()

Samyam Rajbhandari's avatar
Samyam Rajbhandari committed
190
191
start_time = time.time()

Jeff Rasley's avatar
Jeff Rasley committed
192
setup(name='deepspeed',
193
      version=version_str,
Jeff Rasley's avatar
Jeff Rasley committed
194
      description='DeepSpeed library',
195
196
      long_description=readme_text,
      long_description_content_type='text/markdown',
Jeff Rasley's avatar
Jeff Rasley committed
197
198
      author='DeepSpeed Team',
      author_email='deepspeed@microsoft.com',
199
      url='http://deepspeed.ai',
200
      install_requires=install_requires,
201
      extras_require=extras_require,
Jeff Rasley's avatar
Jeff Rasley committed
202
      packages=find_packages(exclude=["docker",
203
204
205
206
207
208
209
                                      "third_party"]),
      include_package_data=True,
      scripts=[
          'bin/deepspeed',
          'bin/deepspeed.pt',
          'bin/ds',
          'bin/ds_ssh',
Jeff Rasley's avatar
Jeff Rasley committed
210
211
          'bin/ds_report',
          'bin/ds_elastic'
212
      ],
213
214
215
216
217
218
      classifiers=[
          'Programming Language :: Python :: 3.6',
          'Programming Language :: Python :: 3.7',
          'Programming Language :: Python :: 3.8'
      ],
      license='MIT',
Jeff Rasley's avatar
Jeff Rasley committed
219
220
      ext_modules=ext_modules,
      cmdclass=cmdclass)
Samyam Rajbhandari's avatar
Samyam Rajbhandari committed
221
222
223

end_time = time.time()
print(f'deepspeed build time = {end_time - start_time} secs')