KungFu/setup_tensorflow.py at master · MachineLearningSystem/KungFu · GitHub

1
2
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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
importmultiprocessing
importos
importsubprocess
importsys
importsysconfig

fromsetuptoolsimportExtension, find_packages, setup
fromsetuptools.command.build_extimportbuild_ext


classCMakeExtension(Extension):
def__init__(self, sourcedir):
Extension.__init__(self, '', sources=[])
self.sourcedir=os.path.abspath(sourcedir)


defensure_absent(filepath):
ifos.path.isfile(filepath):
os.remove(filepath)


defcmake_flag(k, v):
return'-D%s=%s'% (k, str(v))


defcmake_tf_ext_flags():
importtensorflowastf
tf_compile_flags=' '.join(tf.sysconfig.get_compile_flags())
tf_link_flags=' '.join(tf.sysconfig.get_link_flags())
return [
cmake_flag('TF_INCLUDE', tf.sysconfig.get_include()),
cmake_flag('TF_COMPILE_FLAGS', tf_compile_flags),
cmake_flag('TF_LINK_FLAGS', tf_link_flags),
# sysconfig.get_config_var('EXT_SUFFIX') does't work for python2
cmake_flag('PY_EXT_SUFFIX', '%s'%sysconfig.get_config_var('SO')),
 ]


defpass_env(keys):
forkeyinkeys:
val=os.getenv(key)
ifval:
print('Using %s=%s from env'% (key, val))
yieldcmake_flag(key, val)


classCMakeBuild(build_ext):
defbuild_extension(self, ext):
ifnotos.path.exists(self.build_temp):
os.makedirs(self.build_temp)

extdir=self.get_ext_fullpath(ext.name)
ifnotos.path.exists(extdir):
os.makedirs(extdir)

install_prefix=os.path.abspath(os.path.dirname(extdir))
executable_dir=os.path.abspath(os.path.dirname(sys.executable))

# FIXME: do it in a more standard way
if'--user'insys.argv:
executable_dir=os.path.join(os.getenv('HOME'), '.local/bin')

cmake_args= [
# FIXME: use CMAKE_LIBRARY_OUTPUT_DIRECTORY
cmake_flag('LIBRARY_OUTPUT_PATH',
os.path.join(install_prefix, 'kungfu')),
cmake_flag('KUNGFU_BUILD_TF_OPS', 1),
cmake_flag('CMAKE_RUNTIME_OUTPUT_DIRECTORY', executable_dir),
cmake_flag('PYTHON', sys.executable),
 ] +cmake_tf_ext_flags() +list(
pass_env([
'KUNGFU_ENABLE_TRACE',
'CMAKE_VERBOSE_MAKEFILE',
'CMAKE_EXPORT_COMPILE_COMMANDS',
 ]))

ifsys.platform=='linux':
cmake_args.append(cmake_flag('KUNGFU_ENABLE_AFFINITY', 1))
cmake_args.extend(pass_env(['KUNGFU_ENABLE_HWLOC']))

use_nccl=os.getenv('KUNGFU_ENABLE_NCCL')
ifuse_nccl:
cmake_args.append(cmake_flag('KUNGFU_ENABLE_NCCL', use_nccl))
nccl_home=os.getenv('NCCL_HOME')
ifnccl_home:
cmake_args.append(cmake_flag('NCCL_HOME', nccl_home))

ensure_absent(os.path.join(ext.sourcedir, 'CMakeCache.txt'))

subprocess.check_call(
 ['cmake', ext.sourcedir] +cmake_args,
cwd=self.build_temp,
 )
if (os.getenv('CMAKE_BUILD_PARALLEL_LEVEL') isNone
andos.getenv('READTHEDOCS') isNone):
os.environ['CMAKE_BUILD_PARALLEL_LEVEL'] =str(
multiprocessing.cpu_count())
subprocess.check_call(
 ['cmake', '--build', '.'],
cwd=self.build_temp,
 )


package_dir='./srcs/python'

setup(
name='kungfu',
version='0.2.2',
package_dir={'': package_dir},
packages=find_packages(package_dir),
description='KungFu distributed machine learning framework',
url='https://github.com/lsds/KungFu',
ext_modules=[
CMakeExtension('.'),
 ],
cmdclass=dict(build_ext=CMakeBuild),
setup_requires=[],
install_requires=[],
entry_points={
'console_scripts': [
'kungfu-run = kungfu.cmd:run',
 ],
 },
)