Merge pull request #92 from sony/feature/20180907-multi-gpu-package

Feature/20180907 multi gpu package
sony · Sep 10, 2018 · 3044186 · 3044186
2 parents 8702fa1 + fb55451
commit 3044186
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 1 deletion.
diff --git a/build-tools/make/build-with-docker.mk b/build-tools/make/build-with-docker.mk
@@ -94,8 +94,9 @@ bwd-nnabla-ext-cuda-wheel: docker_image_build_cuda
 
 .PHONY: bwd-nnabla-ext-cuda-wheel-multi-gpu
 bwd-nnabla-ext-cuda-wheel-multi-gpu: docker_image_build_cuda_multi_gpu
+	mkdir -p ~/.ccache
 	cd $(NNABLA_EXT_CUDA_DIRECTORY) \
-	&& docker run $(DOCKER_RUN_OPTS) $(DOCKER_IMAGE_BUILD_NNABLA_EXT_CUDA_MULTI_GPU) make -f build-tools/make/build.mk nnabla-ext-cuda-wheel-multi-gpu
+	&& docker run $(DOCKER_RUN_OPTS) $(DOCKER_IMAGE_BUILD_NNABLA_EXT_CUDA_MULTI_GPU) make MULTI_GPU=True -f build-tools/make/build.mk nnabla-ext-cuda-wheel-multi-gpu
 
 .PHONY: bwd-nnabla-ext-cuda-test
 bwd-nnabla-ext-cuda-test: docker_image_build_cuda

diff --git a/build-tools/make/options.mk b/build-tools/make/options.mk
@@ -29,6 +29,9 @@ DOCKER_RUN_OPTS += -e CUDNN_VERSION=$(CUDNN_VERSION)
 export WHL_NO_PREFIX ?= False
 DOCKER_RUN_OPTS += -e WHL_NO_PREFIX=$(WHL_NO_PREFIX)
 
+export MULTI_GPU ?= False
+DOCKER_RUN_OPTS += -e MULTI_GPU=$(MULTI_GPU)
+
 ifndef NNABLA_OPTIONS_INCLUDED
   include $(NNABLA_DIRECTORY)/build-tools/make/options.mk
 endif

diff --git a/python/setup.py b/python/setup.py
@@ -179,6 +179,10 @@ def get_setup_config(root_dir):
     cuda_version = ''
     if 'WHL_NO_PREFIX' in os.environ and os.environ['WHL_NO_PREFIX'] == 'True':
         cuda_version = ''
+    elif 'MULTI_GPU' in os.environ and os.environ['MULTI_GPU'] == 'True':
+        cuda_version = os.environ['CUDA_VERSION_MAJOR'] + \
+            os.environ['CUDA_VERSION_MINOR'] + \
+            '_nccl2_ubuntu16'
     elif 'CUDA_VERSION_MAJOR' in os.environ:
         cuda_version = os.environ['CUDA_VERSION_MAJOR'] + \
             os.environ['CUDA_VERSION_MINOR']