diff --git a/.gitignore b/.gitignore index 12fc56b1c..40f3f6deb 100644 --- a/.gitignore +++ b/.gitignore @@ -141,3 +141,6 @@ docs/.build # pytorch checkpoint *.pt + +# ignore version.py generated by setup.py +colossalai/version.py diff --git a/colossalai/__init__.py b/colossalai/__init__.py index ff65f0f9c..f859161f7 100644 --- a/colossalai/__init__.py +++ b/colossalai/__init__.py @@ -7,4 +7,11 @@ from .initialize import ( launch_from_torch, ) -__version__ = '0.1.11rc4' +try: + # .version will be created by setup.py + from .version import __version__ +except ModuleNotFoundError: + # this will only happen if the user did not run `pip install` + # and directly set PYTHONPATH to use Colossal-AI which is a bad practice + __version__ = '0.0.0' + print('please install Colossal-AI from https://www.colossalai.org/download or from source') diff --git a/requirements/requirements.txt b/requirements/requirements.txt index 528bc6f25..5ac4a3c60 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -1,4 +1,3 @@ -torch>=1.8 numpy tqdm psutil @@ -7,4 +6,4 @@ pre-commit rich click fabric -contexttimer \ No newline at end of file +contexttimer diff --git a/setup.py b/setup.py index 0a83e622e..9766d92f6 100644 --- a/setup.py +++ b/setup.py @@ -4,6 +4,19 @@ import subprocess from setuptools import Extension, find_packages, setup +try: + import torch + from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension + print("\n\ntorch.__version__ = {}\n\n".format(torch.__version__)) + TORCH_MAJOR = int(torch.__version__.split('.')[0]) + TORCH_MINOR = int(torch.__version__.split('.')[1]) + + if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 10): + raise RuntimeError("Colossal-AI requires Pytorch 1.10 or newer.\n" + "The latest stable release can be obtained from https://pytorch.org/") +except ImportError: + raise ModuleNotFoundError('torch is not found. You need to install PyTorch before installing Colossal-AI.') + # ninja build does not work unless include_dirs are abs path this_dir = os.path.dirname(os.path.abspath(__file__)) build_cuda_ext = True @@ -93,29 +106,24 @@ def fetch_readme(): def get_version(): - with open('version.txt') as f: + setup_file_path = os.path.abspath(__file__) + project_path = os.path.dirname(setup_file_path) + version_txt_path = os.path.join(project_path, 'version.txt') + version_py_path = os.path.join(project_path, 'colossalai/version.py') + + with open(version_txt_path) as f: version = f.read().strip() if build_cuda_ext: torch_version = '.'.join(torch.__version__.split('.')[:2]) cuda_version = '.'.join(get_cuda_bare_metal_version(CUDA_HOME)[1:]) version += f'+torch{torch_version}cu{cuda_version}' - return version + # write version into version.py + with open(version_py_path, 'w') as f: + f.write(f"__version__ = '{version}'\n") -if build_cuda_ext: - try: - import torch - from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension - print("\n\ntorch.__version__ = {}\n\n".format(torch.__version__)) - TORCH_MAJOR = int(torch.__version__.split('.')[0]) - TORCH_MINOR = int(torch.__version__.split('.')[1]) + return version - if TORCH_MAJOR < 1 or (TORCH_MAJOR == 1 and TORCH_MINOR < 8): - raise RuntimeError("Colossal-AI requires Pytorch 1.8 or newer.\n" - "The latest stable release can be obtained from https://pytorch.org/") - except ImportError: - print('torch is not found. CUDA extension will not be installed') - build_cuda_ext = False if build_cuda_ext: build_cuda_ext = check_cuda_availability(CUDA_HOME) and check_cuda_torch_binary_vs_bare_metal(CUDA_HOME)