1 from __future__
import absolute_import
2 from __future__
import division
3 from __future__
import print_function
4 from __future__
import unicode_literals
6 from distutils.spawn
import find_executable
7 from distutils
import sysconfig, log
9 import setuptools.command.build_py
10 import setuptools.command.develop
11 import setuptools.command.build_ext
13 from collections
import namedtuple
14 from contextlib
import contextmanager
17 import multiprocessing
21 from textwrap
import dedent
23 TOP_DIR = os.path.realpath(os.path.dirname(__file__))
24 SRC_DIR = os.path.join(TOP_DIR,
'caffe2')
25 CMAKE_BUILD_DIR = os.path.join(TOP_DIR,
'.setuptools-cmake-build')
35 CMAKE = find_executable(
'cmake')
36 assert CMAKE,
'Could not find "cmake" executable!' 37 NINJA = find_executable(
'ninja')
38 MAKE = find_executable(
'make')
39 assert NINJA
or MAKE, \
40 'Could not find neither "ninja" nor "make" executable!' 49 if not os.path.isabs(path):
50 raise RuntimeError(
'Can only cd to absolute path, got: {}'.format(path))
51 orig_path = os.getcwd()
63 git_version = subprocess.check_output([
'git',
'describe',
'--tags',
'HEAD'],
64 cwd=TOP_DIR).decode(
'ascii').strip()
65 except (OSError, subprocess.CalledProcessError):
68 with open(os.path.join(SRC_DIR,
'VERSION_NUMBER'))
as version_file:
69 VersionInfo = namedtuple(
'VersionInfo', [
'version',
'git_version'])(
70 version=version_file.read().strip(),
71 git_version=git_version
82 def initialize_options(self):
85 def finalize_options(self):
89 class create_version(Caffe2Command):
91 with open(os.path.join(SRC_DIR,
'version.py'),
'w')
as f:
94 git_version = '{git_version}' 95 '''.format(**dict(VersionInfo._asdict()))))
100 Compiles everything when `python setup.py build` is run using cmake. 102 Custom args can be passed to cmake by specifying the `CMAKE_ARGS` 103 environment variable. E.g. to build without cuda support run: 104 `CMAKE_ARGS=-DUSE_CUDA=Off python setup.py build` 106 The number of CPUs used by `make`/`ninja` can be specified by passing 107 `-j<ncpus>` to `setup.py build`. By default all CPUs are used. 110 (str(
'jobs='), str(
'j'),
111 str(
'Specifies the number of jobs to use with make or ninja'))
116 def initialize_options(self):
117 self.
jobs = multiprocessing.cpu_count()
119 def finalize_options(self):
123 if cmake_build.built:
125 cmake_build.built =
True 127 if not os.path.exists(CMAKE_BUILD_DIR):
128 os.makedirs(CMAKE_BUILD_DIR)
130 with cd(CMAKE_BUILD_DIR):
133 find_executable(
'cmake'),
134 '-DBUILD_SHARED_LIBS=OFF',
135 '-DPYTHON_EXECUTABLE:FILEPATH={}'.format(sys.executable),
136 '-DPYTHON_INCLUDE_DIR={}'.format(sysconfig.get_python_inc()),
138 '-DBUILD_BENCHMARK=OFF',
139 '-DBUILD_BINARY=OFF',
142 cmake_args.extend([
'-G',
'Ninja'])
143 if 'CMAKE_ARGS' in os.environ:
144 extra_cmake_args = shlex.split(os.environ[
'CMAKE_ARGS'])
146 del os.environ[
'CMAKE_ARGS']
147 log.info(
'Extra cmake args: {}'.format(extra_cmake_args))
148 cmake_args.append(TOP_DIR)
149 subprocess.check_call(cmake_args)
151 build_args = [NINJA
or MAKE]
153 if self.
jobs is not None:
154 build_args.extend([
'-j', str(self.
jobs)])
155 subprocess.check_call(build_args)
158 class build_py(setuptools.command.build_py.build_py):
160 self.run_command(
'create_version')
161 self.run_command(
'cmake_build')
162 for d
in [
'caffe',
'caffe2']:
163 for src
in glob.glob(
164 os.path.join(CMAKE_BUILD_DIR, d,
'proto',
'*.py')):
166 TOP_DIR, os.path.relpath(src, CMAKE_BUILD_DIR))
167 self.copy_file(src, dst)
168 setuptools.command.build_py.build_py.run(self)
171 class build_ext(setuptools.command.build_ext.build_ext):
172 def get_outputs(self):
173 return [os.path.join(self.build_lib, d)
174 for d
in [
'caffe',
'caffe2']]
177 self.run_command(
'cmake_build')
178 setuptools.command.build_ext.build_ext.run(self)
180 def build_extensions(self):
182 while i < len(self.extensions):
183 ext = self.extensions[i]
184 fullname = self.get_ext_fullname(ext.name)
185 filename = self.get_ext_filename(fullname)
187 src = os.path.join(CMAKE_BUILD_DIR, filename)
188 if not os.path.exists(src):
189 del self.extensions[i]
191 dst = os.path.join(os.path.realpath(self.build_lib), filename)
192 self.copy_file(src, dst)
196 class develop(setuptools.command.develop.develop):
198 self.run_command(
'build_py')
199 setuptools.command.develop.develop.run(self)
203 'create_version': create_version,
204 'cmake_build': cmake_build,
205 'build_py': build_py,
206 'build_ext': build_ext,
215 setuptools.Extension(
216 name=str(
'caffe2.python.caffe2_pybind11_state'),
218 setuptools.Extension(
219 name=str(
'caffe2.python.caffe2_pybind11_state_gpu'),
227 packages = setuptools.find_packages()
229 install_requires.extend([
'protobuf',
241 setup_requires.append(
'pytest-runner')
242 tests_require.extend([
'pytest-cov',
'hypothesis'])
250 version=VersionInfo.version,
251 description=
'Caffe2',
252 ext_modules=ext_modules,
255 install_requires=install_requires,
256 setup_requires=setup_requires,
257 tests_require=tests_require,
259 author_email=
'jiayq@fb.com',
260 url=
'https://caffe2.ai',