Add commands to re-build individual libraries. (#2506)
When working on PyTorch dependencies we often want to rebuild only that
dependency and the Python extension. You can now do that by running:
python setup.py build_thc
to only re-build THC
diff --git a/setup.py b/setup.py
index c295a8a..9cc5191 100644
--- a/setup.py
+++ b/setup.py
@@ -77,6 +77,21 @@
# Custom build commands
################################################################################
+dep_libs = [
+ 'TH', 'THS', 'THNN', 'THC', 'THCS', 'THCUNN', 'nccl', 'THPP', 'libshm',
+ 'ATen', 'gloo', 'THD',
+]
+
+
+def build_libs(libs):
+ for lib in libs:
+ assert lib in dep_libs, 'invalid lib: {}'.format(lib)
+ build_libs_cmd = ['bash', 'torch/lib/build_libs.sh']
+ if WITH_CUDA:
+ build_libs_cmd += ['--with-cuda']
+ if subprocess.call(build_libs_cmd + libs) != 0:
+ sys.exit(1)
+
class build_deps(Command):
user_options = []
@@ -88,19 +103,35 @@
pass
def run(self):
- from tools.nnwrap import generate_wrappers as generate_nn_wrappers
- build_all_cmd = ['bash', 'torch/lib/build_all.sh']
+ libs = ['TH', 'THS', 'THNN']
if WITH_CUDA:
- build_all_cmd += ['--with-cuda']
+ libs += ['THC', 'THCS', 'THCUNN']
if WITH_NCCL and not SYSTEM_NCCL:
- build_all_cmd += ['--with-nccl']
+ libs += ['nccl']
+ libs += ['THPP', 'libshm', 'ATen']
if WITH_DISTRIBUTED:
- build_all_cmd += ['--with-distributed']
- if subprocess.call(build_all_cmd) != 0:
- sys.exit(1)
+ if sys.platform == 'linux':
+ libs += ['gloo']
+ libs += ['THD']
+ build_libs(libs)
+
+ from tools.nnwrap import generate_wrappers as generate_nn_wrappers
generate_nn_wrappers()
+build_dep_cmds = {}
+
+for lib in dep_libs:
+ # wrap in function to capture lib
+ class build_dep(build_deps):
+ description = 'Build {} external library'.format(lib)
+
+ def run(self):
+ build_libs([self.lib])
+ build_dep.lib = lib
+ build_dep_cmds['build_' + lib.lower()] = build_dep
+
+
class build_module(Command):
user_options = []
@@ -470,26 +501,30 @@
except subprocess.CalledProcessError:
pass
+cmdclass = {
+ 'build': build,
+ 'build_py': build_py,
+ 'build_ext': build_ext,
+ 'build_deps': build_deps,
+ 'build_module': build_module,
+ 'develop': develop,
+ 'install': install,
+ 'clean': clean,
+}
+cmdclass.update(build_dep_cmds)
setup(name="torch", version=version,
description="Tensors and Dynamic neural networks in Python with strong GPU acceleration",
ext_modules=extensions,
- cmdclass={
- 'build': build,
- 'build_py': build_py,
- 'build_ext': build_ext,
- 'build_deps': build_deps,
- 'build_module': build_module,
- 'develop': develop,
- 'install': install,
- 'clean': clean,
- },
+ cmdclass=cmdclass,
packages=packages,
package_data={'torch': [
'lib/*.so*', 'lib/*.dylib*',
'lib/torch_shm_manager',
'lib/*.h',
'lib/include/TH/*.h', 'lib/include/TH/generic/*.h',
- 'lib/include/THC/*.h', 'lib/include/THC/generic/*.h']},
+ 'lib/include/THC/*.h', 'lib/include/THC/generic/*.h',
+ 'lib/include/ATen/*.h',
+ ]},
install_requires=['pyyaml', 'numpy'],
)
diff --git a/torch/lib/build_all.sh b/torch/lib/build_libs.sh
similarity index 79%
rename from torch/lib/build_all.sh
rename to torch/lib/build_libs.sh
index bcd2799..4594b0c 100755
--- a/torch/lib/build_all.sh
+++ b/torch/lib/build_libs.sh
@@ -10,19 +10,10 @@
# Options for building only a subset of the libraries
WITH_CUDA=0
-WITH_NCCL=0
-WITH_DISTRIBUTED=0
-for arg in "$@"; do
- if [[ "$arg" == "--with-cuda" ]]; then
- WITH_CUDA=1
- elif [[ "$arg" == "--with-nccl" ]]; then
- WITH_NCCL=1
- elif [[ "$arg" == "--with-distributed" ]]; then
- WITH_DISTRIBUTED=1
- else
- echo "Unknown argument: $arg"
- fi
-done
+if [[ "$1" == "--with-cuda" ]]; then
+ WITH_CUDA=1
+ shift
+fi
cd "$(dirname "$0")/../.."
BASE_DIR=$(pwd)
@@ -43,6 +34,11 @@
else
LDFLAGS="$LDFLAGS -Wl,-rpath,\$ORIGIN"
fi
+CPP_FLAGS=" -std=c++11 "
+GLOO_FLAGS=""
+if [[ $WITH_CUDA -eq 1 ]]; then
+ GLOO_FLAGS="-DUSE_CUDA=1 -DNCCL_ROOT_DIR=$INSTALL_DIR"
+fi
# Used to build an individual library, e.g. build TH
function build() {
@@ -98,6 +94,7 @@
cd ../..
fi
}
+
function build_nccl() {
mkdir -p build/nccl
cd build/nccl
@@ -117,43 +114,16 @@
# In the torch/lib directory, create an installation directory
mkdir -p tmp_install
-# We need to build the CPU libraries first, because they are used
-# in the CUDA libraries
-build TH
-build THS
-build THNN
-
-CPP_FLAGS=" -std=c++11 "
-if [[ $WITH_CUDA -eq 1 ]]; then
- build THC
- build THCS
- build THCUNN
-fi
-if [[ $WITH_NCCL -eq 1 ]]; then
- build_nccl
-fi
-
-# THPP has dependencies on both CPU and CUDA, so build it
-# after those libraries have been completed
-build THPP
-
-# The shared memory manager depends on TH
-build libshm
-build ATen
-
-# THD, gloo have dependencies on Torch, CUDA, NCCL etc.
-if [[ $WITH_DISTRIBUTED -eq 1 ]]; then
- if [ "$(uname)" == "Linux" ]; then
- if [ -d "gloo" ]; then
- GLOO_FLAGS=""
- if [[ $WITH_CUDA -eq 1 ]]; then
- GLOO_FLAGS="-DUSE_CUDA=1 -DNCCL_ROOT_DIR=$INSTALL_DIR"
- fi
- build gloo "$GLOO_FLAGS"
- fi
+# Build
+for arg in "$@"; do
+ if [[ "$arg" == "nccl" ]]; then
+ build_nccl
+ elif [[ "$arg" == "gloo" ]]; then
+ build gloo $GLOO_FLAGS
+ else
+ build $arg
fi
- build THD
-fi
+done
# If all the builds succeed we copy the libraries, headers,
# binaries to torch/lib
@@ -161,7 +131,9 @@
cp THNN/generic/THNN.h .
cp THCUNN/generic/THCUNN.h .
cp -r $INSTALL_DIR/include .
-cp $INSTALL_DIR/bin/* .
+if [ -d "$INSTALL_DIR/bin/" ]; then
+ cp $INSTALL_DIR/bin/* .
+fi
# this is for binary builds
if [[ $PYTORCH_BINARY_BUILD && $PYTORCH_SO_DEPS ]]
