Caffe2 - Python API
A deep learning, cross platform ML framework
setup.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 assert find_executable('cmake'), 'Could not find "cmake" executable!'
36 assert find_executable('make'), 'Could not find "make" executable!'
37 
38 ################################################################################
39 # utils functions
40 ################################################################################
41 
42 
43 @contextmanager
44 def cd(path):
45  if not os.path.isabs(path):
46  raise RuntimeError('Can only cd to absolute path, got: {}'.format(path))
47  orig_path = os.getcwd()
48  os.chdir(path)
49  try:
50  yield
51  finally:
52  os.chdir(orig_path)
53 
54 ################################################################################
55 # Version
56 ################################################################################
57 
58 try:
59  git_version = subprocess.check_output(['git', 'describe', '--tags', 'HEAD'],
60  cwd=TOP_DIR).decode('ascii').strip()
61 except (OSError, subprocess.CalledProcessError):
62  git_version = None
63 
64 with open(os.path.join(TOP_DIR, 'VERSION_NUMBER')) as version_file:
65  VersionInfo = namedtuple('VersionInfo', ['version', 'git_version'])(
66  version=version_file.read().strip(),
67  git_version=git_version
68  )
69 
70 ################################################################################
71 # Customized commands
72 ################################################################################
73 
74 
75 class Caffe2Command(setuptools.Command):
76  user_options = []
77 
78  def initialize_options(self):
79  pass
80 
81  def finalize_options(self):
82  pass
83 
84 
85 class create_version(Caffe2Command):
86  def run(self):
87  with open(os.path.join(SRC_DIR, 'version.py'), 'w') as f:
88  f.write(dedent('''
89  version = '{version}'
90  git_version = '{git_version}'
91  '''.format(**dict(VersionInfo._asdict()))))
92 
93 
95  """
96  Compiles everything when `python setup.py build` is run using cmake.
97 
98  Custom args can be passed to cmake by specifying the `CMAKE_ARGS`
99  environment variable. E.g. to build without cuda support run:
100  `CMAKE_ARGS=-DUSE_CUDA=Off python setup.py build`
101 
102  The number of CPUs used by `make` can be specified by passing `-j<ncpus>`
103  to `setup.py build`. By default all CPUs are used.
104  """
105  user_options = [
106  (str('jobs='), str('j'), str('Specifies the number of jobs to use with make'))
107  ]
108 
109  built = False
110 
111  def initialize_options(self):
112  self.jobs = multiprocessing.cpu_count()
113 
114  def finalize_options(self):
115  self.jobs = int(self.jobs)
116 
117  def run(self):
118  if cmake_build.built:
119  return
120  cmake_build.built = True
121 
122  if not os.path.exists(CMAKE_BUILD_DIR):
123  os.makedirs(CMAKE_BUILD_DIR)
124 
125  with cd(CMAKE_BUILD_DIR):
126  # configure
127  cmake_args = [
128  find_executable('cmake'),
129  '-DBUILD_SHARED_LIBS=OFF',
130  '-DPYTHON_EXECUTABLE:FILEPATH={}'.format(sys.executable),
131  '-DPYTHON_INCLUDE_DIR={}'.format(sysconfig.get_python_inc()),
132  '-DBUILD_TEST=OFF',
133  '-DBUILD_BENCHMARK=OFF',
134  '-DBUILD_BINARY=OFF',
135  ]
136  if 'CMAKE_ARGS' in os.environ:
137  extra_cmake_args = shlex.split(os.environ['CMAKE_ARGS'])
138  # prevent crossfire with downstream scripts
139  del os.environ['CMAKE_ARGS']
140  log.info('Extra cmake args: {}'.format(extra_cmake_args))
141  cmake_args.append(TOP_DIR)
142  subprocess.check_call(cmake_args)
143 
144  build_args = [find_executable('make')]
145  # control the number of concurrent jobs
146  if self.jobs is not None:
147  build_args.extend(['-j', str(self.jobs)])
148  subprocess.check_call(build_args)
149 
150 
151 class build_py(setuptools.command.build_py.build_py):
152  def run(self):
153  self.run_command('create_version')
154  self.run_command('cmake_build')
155  for d in ['caffe', 'caffe2']:
156  for src in glob.glob(
157  os.path.join(CMAKE_BUILD_DIR, d, 'proto', '*.py')):
158  dst = os.path.join(
159  TOP_DIR, os.path.relpath(src, CMAKE_BUILD_DIR))
160  self.copy_file(src, dst)
161  setuptools.command.build_py.build_py.run(self)
162 
163 
164 class build_ext(setuptools.command.build_ext.build_ext):
165  def get_outputs(self):
166  return [os.path.join(self.build_lib, d)
167  for d in ['caffe', 'caffe2']]
168 
169  def run(self):
170  self.run_command('cmake_build')
171  setuptools.command.build_ext.build_ext.run(self)
172 
173  def build_extensions(self):
174  i = 0
175  while i < len(self.extensions):
176  ext = self.extensions[i]
177  fullname = self.get_ext_fullname(ext.name)
178  filename = self.get_ext_filename(fullname)
179 
180  src = os.path.join(CMAKE_BUILD_DIR, filename)
181  if not os.path.exists(src):
182  del self.extensions[i]
183  else:
184  dst = os.path.join(os.path.realpath(self.build_lib), filename)
185  self.copy_file(src, dst)
186  i += 1
187 
188 
189 class develop(setuptools.command.develop.develop):
190  def run(self):
191  self.run_command('build_py')
192  setuptools.command.develop.develop.run(self)
193 
194 
195 cmdclass = {
196  'create_version': create_version,
197  'cmake_build': cmake_build,
198  'build_py': build_py,
199  'build_ext': build_ext,
200  'develop': develop,
201 }
202 
203 ################################################################################
204 # Extensions
205 ################################################################################
206 
207 ext_modules = [
208  setuptools.Extension(
209  name=str('caffe2.python.caffe2_pybind11_state'),
210  sources=[]),
211  setuptools.Extension(
212  name=str('caffe2.python.caffe2_pybind11_state_gpu'),
213  sources=[]),
214 ]
215 
216 ################################################################################
217 # Packages
218 ################################################################################
219 
220 packages = setuptools.find_packages()
221 
222 install_requires.extend(['protobuf',
223  'numpy',
224  'flask',
225  'future',
226  'graphviz',
227  'hypothesis',
228  'jupyter',
229  'matplotlib',
230  'pydot',
231  'python-nvd3',
232  'pyyaml',
233  'requests',
234  'scikit-image',
235  'scipy',
236  'setuptools',
237  'six',
238  'tornado'])
239 
240 ################################################################################
241 # Test
242 ################################################################################
243 
244 setup_requires.append('pytest-runner')
245 tests_require.extend(['pytest-cov', 'hypothesis'])
246 
247 ################################################################################
248 # Final
249 ################################################################################
250 
251 setuptools.setup(
252  name='caffe2',
253  version=VersionInfo.version,
254  description='Caffe2',
255  ext_modules=ext_modules,
256  cmdclass=cmdclass,
257  packages=packages,
258  install_requires=install_requires,
259  setup_requires=setup_requires,
260  tests_require=tests_require,
261  author='jiayq',
262  author_email='[email protected]',
263  url='https://caffe2.ai',
264 )
Customized commands.
Definition: setup.py:75