colossalai

1
import os
2
import sys
3
from typing import List
4

5
from setuptools import find_packages, setup
6

7
try:
8
    import torch  # noqa
9
    from torch.utils.cpp_extension import BuildExtension
10

11
    TORCH_AVAILABLE = True
12
except ImportError:
13
    TORCH_AVAILABLE = False
14

15
THIS_DIR = os.path.dirname(os.path.abspath(__file__))
16
BUILD_EXT = int(os.environ.get("BUILD_EXT", "0")) == 1
17

18
# we do not support windows currently
19
if sys.platform == "win32":
20
    raise RuntimeError("Windows is not supported yet. Please try again within the Windows Subsystem for Linux (WSL).")
21

22

23
def fetch_requirements(path) -> List[str]:
24
    """
25
    This function reads the requirements file.
26

27
    Args:
28
        path (str): the path to the requirements file.
29

30
    Returns:
31
        The lines in the requirements file.
32
    """
33
    with open(path, "r") as fd:
34
        return [r.strip() for r in fd.readlines()]
35

36

37
def fetch_readme() -> str:
38
    """
39
    This function reads the README.md file in the current directory.
40

41
    Returns:
42
        The lines in the README file.
43
    """
44
    with open("README.md", encoding="utf-8") as f:
45
        return f.read()
46

47

48
def get_version() -> str:
49
    """
50
    This function reads the version.txt and generates the colossalai/version.py file.
51

52
    Returns:
53
        The library version stored in version.txt.
54
    """
55

56
    setup_file_path = os.path.abspath(__file__)
57
    project_path = os.path.dirname(setup_file_path)
58
    version_txt_path = os.path.join(project_path, "version.txt")
59
    version_py_path = os.path.join(project_path, "colossalai/version.py")
60

61
    with open(version_txt_path) as f:
62
        version = f.read().strip()
63

64
    # write version into version.py
65
    with open(version_py_path, "w") as f:
66
        f.write(f"__version__ = '{version}'\n")
67
    return version
68

69

70
if BUILD_EXT:
71
    if not TORCH_AVAILABLE:
72
        raise ModuleNotFoundError(
73
            "[extension] PyTorch is not found while BUILD_EXT=1. You need to install PyTorch first in order to build CUDA extensions"
74
        )
75

76
    from extensions import ALL_EXTENSIONS
77

78
    op_names = []
79
    ext_modules = []
80

81
    for ext_cls in ALL_EXTENSIONS:
82
        ext = ext_cls()
83
        if ext.support_aot and ext.is_hardware_available():
84
            ext.assert_hardware_compatible()
85
            op_names.append(ext.name)
86
            ext_modules.append(ext.build_aot())
87

88
    # show log
89
    if len(ext_modules) == 0:
90
        raise RuntimeError("[extension] Could not find any kernel compatible with the current environment.")
91
    else:
92
        op_name_list = ", ".join(op_names)
93
        print(f"[extension] Building extensions{op_name_list}")
94
else:
95
    ext_modules = []
96

97
version = get_version()
98
package_name = "colossalai"
99

100
setup(
101
    name=package_name,
102
    version=version,
103
    packages=find_packages(
104
        exclude=(
105
            "extensions",
106
            "benchmark",
107
            "docker",
108
            "tests",
109
            "docs",
110
            "examples",
111
            "tests",
112
            "scripts",
113
            "requirements",
114
            "extensions",
115
            "*.egg-info",
116
        ),
117
    ),
118
    description="An integrated large-scale model training system with efficient parallelization techniques",
119
    long_description=fetch_readme(),
120
    long_description_content_type="text/markdown",
121
    license="Apache Software License 2.0",
122
    url="https://www.colossalai.org",
123
    project_urls={
124
        "Forum": "https://github.com/hpcaitech/ColossalAI/discussions",
125
        "Bug Tracker": "https://github.com/hpcaitech/ColossalAI/issues",
126
        "Examples": "https://github.com/hpcaitech/ColossalAI-Examples",
127
        "Documentation": "http://colossalai.readthedocs.io",
128
        "Github": "https://github.com/hpcaitech/ColossalAI",
129
    },
130
    ext_modules=ext_modules,
131
    cmdclass={"build_ext": BuildExtension} if ext_modules else {},
132
    install_requires=fetch_requirements("requirements/requirements.txt"),
133
    entry_points="""
134
        [console_scripts]
135
        colossalai=colossalai.cli:cli
136
    """,
137
    python_requires=">=3.6",
138
    classifiers=[
139
        "Programming Language :: Python :: 3",
140
        "License :: OSI Approved :: Apache Software License",
141
        "Environment :: GPU :: NVIDIA CUDA",
142
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
143
        "Topic :: System :: Distributed Computing",
144
    ],
145
    package_data={
146
        "colossalai": [
147
            "kernel/extensions/csrc/**/*",
148
        ]
149
    },
150
)
151