setup.py 7.9 KB
Newer Older
liuzhe-lz's avatar
liuzhe-lz committed
1
2
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
3

4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
"""
Script for installation and distribution.

You can use environment variable `NNI_RELEASE` to set release version.

If release version is not set, default to a development build whose version string will be `999.dev0`.


## Development ##

Build and install for development:

  $ python setup.py develop

Uninstall:

  $ pip uninstall nni

Remove generated files: (use "--all" to remove toolchain and built wheel)

  $ python setup.py clean [--all]

Build TypeScript modules without install:

  $ python setup.py build_ts


## Release ##

Build wheel package:

  $ NNI_RELEASE=2.0 python setup.py build_ts
  $ NNI_RELEASE=2.0 python setup.py bdist_wheel -p manylinux1_x86_64

Where "2.0" is version string and "manylinux1_x86_64" is platform.
The platform may also be "macosx_10_9_x86_64" or "win_amd64".

`build_ts` must be manually invoked before `bdist_wheel`,
or setuptools cannot locate JS files which should be packed into wheel.
"""

from distutils.cmd import Command
from distutils.command.build import build
from distutils.command.clean import clean
import glob
49
import os
50
import shutil
liuzhe-lz's avatar
liuzhe-lz committed
51
import sys
52
53
54
55
56

import setuptools
from setuptools.command.develop import develop

import setup_ts
57
58


59
60
61
62
63
64
65
66
67
68
69
70
71
72
dependencies = [
    'astor',
    'hyperopt==0.1.2',
    'json_tricks',
    'netifaces',
    'psutil',
    'ruamel.yaml',
    'requests',
    'responses',
    'schema',
    'PythonWebHDFS',
    'colorama',
    'scikit-learn>=0.23.2',
    'pkginfo',
colorjam's avatar
colorjam committed
73
    'websockets',
74
    'filelock',
liuzhe-lz's avatar
liuzhe-lz committed
75
    'prettytable',
76
    'dataclasses ; python_version < "3.7"',
liuzhe-lz's avatar
liuzhe-lz committed
77
78
    'numpy < 1.19.4 ; sys_platform == "win32"',
    'numpy < 1.20 ; sys_platform != "win32" and python_version < "3.7"',
liuzhe-lz's avatar
liuzhe-lz committed
79
80
81
    'numpy ; sys.platform != "win32" and python_version >= "3.7"',
    'scipy < 1.6 ; python_version < "3.7"',
    'scipy ; python_version >= "3.7"',
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
]

release = os.environ.get('NNI_RELEASE')

def _setup():
    setuptools.setup(
        name = 'nni',
        version = release or '999.dev0',
        description = 'Neural Network Intelligence project',
        long_description = open('README.md', encoding='utf-8').read(),
        long_description_content_type = 'text/markdown',
        url = 'https://github.com/Microsoft/nni',
        author = 'Microsoft NNI Team',
        author_email = 'nni@microsoft.com',
        license = 'MIT',
        classifiers = [
            'License :: OSI Approved :: MIT License',
            'Operating System :: MacOS :: MacOS X',
            'Operating System :: Microsoft :: Windows :: Windows 10',
            'Operating System :: POSIX :: Linux',
            'Programming Language :: Python :: 3 :: Only',
            'Topic :: Scientific/Engineering :: Artificial Intelligence',
        ],

        packages = _find_python_packages(),
        package_data = {
liuzhe-lz's avatar
liuzhe-lz committed
108
            'nni': _find_requirements_txt(),  # must do this manually due to setuptools issue #1806
109
110
111
112
113
            'nni_node': _find_node_files()  # note: this does not work before building
        },

        python_requires = '>=3.6',
        install_requires = dependencies,
114
115
116
117
118
119
120
121
        extras_require = {
            'SMAC': [
                'ConfigSpaceNNI @ git+https://github.com/QuanluZhang/ConfigSpace.git',
                'smac @ git+https://github.com/QuanluZhang/SMAC3.git'
            ],
            'BOHB': ['ConfigSpace==0.4.7', 'statsmodels==0.10.0'],
            'PPOTuner': ['enum34', 'gym']
        },
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
        setup_requires = ['requests'],

        entry_points = {
            'console_scripts' : [
                'nnictl = nni.tools.nnictl.nnictl:parse_args'
            ]
        },

        cmdclass = {
            'build': Build,
            'build_ts': BuildTs,
            'clean': Clean,
            'develop': Develop,
        }
    )
137
138
139
140
141


def _find_python_packages():
    packages = []
    for dirpath, dirnames, filenames in os.walk('nni'):
liuzhe-lz's avatar
liuzhe-lz committed
142
        if '/__pycache__' not in dirpath and '/.mypy_cache' not in dirpath:
143
144
145
            packages.append(dirpath.replace('/', '.'))
    return sorted(packages) + ['nni_node']

liuzhe-lz's avatar
liuzhe-lz committed
146
147
148
149
150
151
152
def _find_requirements_txt():
    requirement_files = []
    for dirpath, dirnames, filenames in os.walk('nni'):
        if 'requirements.txt' in filenames:
            requirement_files.append(os.path.join(dirpath[len('nni/'):], 'requirements.txt'))
    return requirement_files

153
def _find_node_files():
154
    if not os.path.exists('nni_node'):
liuzhe-lz's avatar
liuzhe-lz committed
155
156
        if release and 'build_ts' not in sys.argv:
            sys.exit('ERROR: To build a release version, run "python setup.py build_ts" first')
157
        return []
158
159
160
    files = []
    for dirpath, dirnames, filenames in os.walk('nni_node'):
        for filename in filenames:
liuzhe-lz's avatar
liuzhe-lz committed
161
            files.append(os.path.join(dirpath[len('nni_node/'):], filename))
162
163
    if '__init__.py' in files:
        files.remove('__init__.py')
164
    return sorted(files)
165

liuzhe-lz's avatar
liuzhe-lz committed
166
167
168
def _using_conda_or_virtual_environment():
    return sys.prefix != sys.base_prefix or os.path.isdir(os.path.join(sys.prefix, 'conda-meta'))

169
170
171
172
173
174
175
176
177
178
179
180
181
def _copy_data_files():
    # after installation, nni needs to find this location in nni.tools.package_utils.get_registered_algo_config_path
    # since we can not import nni here, we need to ensure get_registered_algo_config_path use the same
    # logic here to retrieve registered_algorithms.yml
    if _using_conda_or_virtual_environment():
        nni_config_dir = os.path.join(sys.prefix, 'nni')
    elif sys.platform == 'win32':
        nni_config_dir = os.path.join(os.getenv('APPDATA'), 'nni')
    else:
        nni_config_dir = os.path.expanduser('~/.config/nni')
    if not os.path.exists(nni_config_dir):
        os.makedirs(nni_config_dir)
    shutil.copyfile('./deployment/registered_algorithms.yml', os.path.join(nni_config_dir, 'registered_algorithms.yml'))
182

183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
class BuildTs(Command):
    description = 'build TypeScript modules'

    user_options = []

    def initialize_options(self):
        pass

    def finalize_options(self):
        pass

    def run(self):
        setup_ts.build(release)

class Build(build):
    def run(self):
liuzhe-lz's avatar
liuzhe-lz committed
199
200
201
202
        if not release:
            sys.exit('Please set environment variable "NNI_RELEASE=<release_version>"')
        if os.path.islink('nni_node/main.js'):
            sys.exit('A development build already exists. Please uninstall NNI and run "python3 setup.py clean --all".')
203
        _copy_data_files()
204
205
206
        super().run()

class Develop(develop):
liuzhe-lz's avatar
liuzhe-lz committed
207
    user_options = develop.user_options + [
liuzhe-lz's avatar
liuzhe-lz committed
208
209
        ('no-user', None, 'Prevent automatically adding "--user"'),
        ('skip-ts', None, 'Prevent building TypeScript modules')
liuzhe-lz's avatar
liuzhe-lz committed
210
211
    ]

liuzhe-lz's avatar
liuzhe-lz committed
212
    boolean_options = develop.boolean_options + ['no-user', 'skip-ts']
liuzhe-lz's avatar
liuzhe-lz committed
213
214
215
216

    def initialize_options(self):
        super().initialize_options()
        self.no_user = None
liuzhe-lz's avatar
liuzhe-lz committed
217
        self.skip_ts = None
liuzhe-lz's avatar
liuzhe-lz committed
218

219
    def finalize_options(self):
liuzhe-lz's avatar
liuzhe-lz committed
220
221
222
223
        # if `--user` or `--no-user` is explicitly set, do nothing
        # otherwise activate `--user` if using system python
        if not self.user and not self.no_user:
            self.user = not _using_conda_or_virtual_environment()
224
225
226
        super().finalize_options()

    def run(self):
liuzhe-lz's avatar
liuzhe-lz committed
227
228
        if not self.skip_ts:
            setup_ts.build(release=None)
229
        _copy_data_files()
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
        super().run()

class Clean(clean):
    def finalize_options(self):
        self._all = self.all
        self.all = True  # always use `clean --all`
        super().finalize_options()

    def run(self):
        super().run()
        setup_ts.clean(self._all)
        _clean_temp_files()
        shutil.rmtree('nni.egg-info', ignore_errors=True)
        if self._all:
            shutil.rmtree('dist', ignore_errors=True)


def _clean_temp_files():
    for pattern in _temp_files:
        for path in glob.glob(pattern):
            if os.path.islink(path) or os.path.isfile(path):
                os.remove(path)
            else:
                shutil.rmtree(path)

_temp_files = [
    # unit test
    'test/model_path/',
    'test/temp.json',
259
260
    'test/ut/sdk/*.pth',
    'test/ut/tools/annotation/_generated/'
261
262
263
]


liuzhe-lz's avatar
liuzhe-lz committed
264
265
if __name__ == '__main__':
    _setup()