diff --git a/build-tools/make/build-with-docker.mk b/build-tools/make/build-with-docker.mk index 0fb4dc1c2..4b622f653 100644 --- 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 index 1f241c6f7..d809f8159 100644 --- 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 index 65399f96f..54dad29ba 100644 --- 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']