Caffe2 - Python API
A deep learning, cross platform ML framework
setup_caffe2.py
1 from __future__ import absolute_import
2 from __future__ import division
3 from __future__ import print_function
4 from __future__ import unicode_literals
5 
6 from distutils.spawn import find_executable
7 from distutils import sysconfig, log
8 import setuptools
9 import setuptools.command.build_py
10 import setuptools.command.develop
11 import setuptools.command.build_ext
12 
13 from collections import namedtuple
14 from contextlib import contextmanager
15 import glob
16 import os
17 import multiprocessing
18 import shlex
19 import subprocess
20 import sys
21 from textwrap import dedent
22 
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')
26 
27 install_requires = []
28 setup_requires = []
29 tests_require = []
30 
31 ################################################################################
32 # Pre Check
33 ################################################################################
34 
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!'
41 
42 ################################################################################
43 # utils functions
44 ################################################################################
45 
46 
47 @contextmanager
48 def cd(path):
49  if not os.path.isabs(path):
50  raise RuntimeError('Can only cd to absolute path, got: {}'.format(path))
51  orig_path = os.getcwd()
52  os.chdir(path)
53  try:
54  yield
55  finally:
56  os.chdir(orig_path)
57 
58 ################################################################################
59 # Version
60 ################################################################################
61 
62 try:
63  git_version = subprocess.check_output(['git', 'describe', '--tags', 'HEAD'],
64  cwd=TOP_DIR).decode('ascii').strip()
65 except (OSError, subprocess.CalledProcessError):
66  git_version = None
67 
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
72  )
73 
74 ################################################################################
75 # Customized commands
76 ################################################################################
77 
78 
79 class Caffe2Command(setuptools.Command):
80  user_options = []
81 
82  def initialize_options(self):
83  pass
84 
85  def finalize_options(self):
86  pass
87 
88 
89 class create_version(Caffe2Command):
90  def run(self):
91  with open(os.path.join(SRC_DIR, 'version.py'), 'w') as f:
92  f.write(dedent('''
93  version = '{version}'
94  git_version = '{git_version}'
95  '''.format(**dict(VersionInfo._asdict()))))
96 
97 
99  """
100  Compiles everything when `python setup.py build` is run using cmake.
101 
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`
105 
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.
108  """
109  user_options = [
110  (str('jobs='), str('j'),
111  str('Specifies the number of jobs to use with make or ninja'))
112  ]
113 
114  built = False
115 
116  def initialize_options(self):
117  self.jobs = multiprocessing.cpu_count()
118 
119  def finalize_options(self):
120  self.jobs = int(self.jobs)
121 
122  def run(self):
123  if cmake_build.built:
124  return
125  cmake_build.built = True
126 
127  if not os.path.exists(CMAKE_BUILD_DIR):
128  os.makedirs(CMAKE_BUILD_DIR)
129 
130  with cd(CMAKE_BUILD_DIR):
131  # configure
132  cmake_args = [
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()),
137  '-DBUILD_TEST=OFF',
138  '-DBUILD_BENCHMARK=OFF',
139  '-DBUILD_BINARY=OFF',
140  ]
141  if NINJA:
142  cmake_args.extend(['-G', 'Ninja'])
143  if 'CMAKE_ARGS' in os.environ:
144  extra_cmake_args = shlex.split(os.environ['CMAKE_ARGS'])
145  # prevent crossfire with downstream scripts
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)
150 
151  build_args = [NINJA or MAKE]
152  # control the number of concurrent jobs
153  if self.jobs is not None:
154  build_args.extend(['-j', str(self.jobs)])
155  subprocess.check_call(build_args)
156 
157 
158 class build_py(setuptools.command.build_py.build_py):
159  def run(self):
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')):
165  dst = os.path.join(
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)
169 
170 
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']]
175 
176  def run(self):
177  self.run_command('cmake_build')
178  setuptools.command.build_ext.build_ext.run(self)
179 
180  def build_extensions(self):
181  i = 0
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)
186 
187  src = os.path.join(CMAKE_BUILD_DIR, filename)
188  if not os.path.exists(src):
189  del self.extensions[i]
190  else:
191  dst = os.path.join(os.path.realpath(self.build_lib), filename)
192  self.copy_file(src, dst)
193  i += 1
194 
195 
196 class develop(setuptools.command.develop.develop):
197  def run(self):
198  self.run_command('build_py')
199  setuptools.command.develop.develop.run(self)
200 
201 
202 cmdclass = {
203  'create_version': create_version,
204  'cmake_build': cmake_build,
205  'build_py': build_py,
206  'build_ext': build_ext,
207  'develop': develop,
208 }
209 
210 ################################################################################
211 # Extensions
212 ################################################################################
213 
214 ext_modules = [
215  setuptools.Extension(
216  name=str('caffe2.python.caffe2_pybind11_state'),
217  sources=[]),
218  setuptools.Extension(
219  name=str('caffe2.python.caffe2_pybind11_state_gpu'),
220  sources=[]),
221 ]
222 
223 ################################################################################
224 # Packages
225 ################################################################################
226 
227 packages = setuptools.find_packages()
228 
229 install_requires.extend(['protobuf',
230  'numpy',
231  'future',
232  'hypothesis',
233  'requests',
234  'scipy',
235  'six',])
236 
237 ################################################################################
238 # Test
239 ################################################################################
240 
241 setup_requires.append('pytest-runner')
242 tests_require.extend(['pytest-cov', 'hypothesis'])
243 
244 ################################################################################
245 # Final
246 ################################################################################
247 
248 setuptools.setup(
249  name='caffe2',
250  version=VersionInfo.version,
251  description='Caffe2',
252  ext_modules=ext_modules,
253  cmdclass=cmdclass,
254  packages=packages,
255  install_requires=install_requires,
256  setup_requires=setup_requires,
257  tests_require=tests_require,
258  author='jiayq',
259  author_email='[email protected]',
260  url='https://caffe2.ai',
261 )
Customized commands.
Definition: setup_caffe2.py:79