Build distributed libs in build_libtorch.py (#14037)

Summary:
This patch detects and builds c10d and gloo for the C++ API.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/14037

Reviewed By: ezyang

Differential Revision: D13283801

Pulled By: ebetica

fbshipit-source-id: 006dbb691344819833da6b4b844c1f0572942135

diff --git a/tools/build_libtorch.py b/tools/build_libtorch.py
index 8b6f030..6122bea 100644 (file)
--- a/tools/build_libtorch.py
+++ b/tools/build_libtorch.py
@@ -5,6 +5,7 @@ import subprocess
 import sys
 
 from setup_helpers.cuda import USE_CUDA
+from setup_helpers.dist_check import USE_DISTRIBUTED, USE_GLOO_IBVERBS, IS_LINUX
 from setup_helpers.env import check_env_flag
 
 if __name__ == '__main__':
@@ -31,6 +32,11 @@ if __name__ == '__main__':
         command.append('--use-cuda')
         if os.environ.get('USE_CUDA_STATIC_LINK', False):
             command.append('--cuda-static-link')
+    if USE_DISTRIBUTED and IS_LINUX:
+        if USE_GLOO_IBVERBS:
+            command.append('--use-gloo-ibverbs')
+        command.append('--use-distributed')
+
     command.append('caffe2')
 
     sys.stdout.flush()

diff --git a/tools/build_pytorch_libs.sh b/tools/build_pytorch_libs.sh
index 0921bd7..11d5ef5 100755 (executable)
--- a/tools/build_pytorch_libs.sh
+++ b/tools/build_pytorch_libs.sh
@@ -59,6 +59,9 @@ while [[ $# -gt 0 ]]; do
       --use-cuda)
           USE_CUDA=1
           ;;
+      --use-distributed)
+          USE_DISTRIBUTED=1
+          ;;
       --use-fbgemm)
           USE_FBGEMM=1
           ;;