setup.py

#!/usr/bin/env python
import os
import subprocess

from setuptools import find_packages, setup

DEFAULT_VERSION = '2.3.0'


def get_and_set_version():
    version = os.getenv('TORCHACC_VERSION', DEFAULT_VERSION)
    commit_id = subprocess.check_output(['git', 'rev-parse',
                                         'HEAD']).strip().decode('utf-8')[:6]
    py_version_path = os.path.join(
        os.path.dirname(os.path.abspath(__file__)), 'torchacc', 'version.py')
    with open(py_version_path, 'w') as f:
        f.write('# Autogenerated file, do not edit!\n')
        f.write("__version__ = '{}-{}'\n".format(version, commit_id))
    return version


setup_info = dict(
    name=os.environ.get('TORCH_ACC_PACKAGE_NAME', 'torchacc'),
    version=get_and_set_version(),
    author='Alibaba Inc.',
    url='http://gitlab.alibaba-inc.com/torchx/torchacc',
    description='Torch Accelerator powered by Alibaba',
    long_description=open('README.md').read(),
    long_description_content_type='text/markdown',

    # Package info
    packages=['torchacc'] + ['torchacc.' + \
                             pkg for pkg in find_packages('torchacc')],

    # add console_scripts
    entry_points={
        'console_scripts': [
            'consolidate_and_reshard_fsdp_ckpts = torchacc.utils.consolidate_and_reshard_ckpts:main',
        ],
    },

    # Add _ prefix to the names of temporary build dirs
    options={'build': {'build_base': '_build'}, },
    zip_safe=True,
)

setup(**setup_info)