feat: build manylinux compatible builds (#1527)

Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
pytorch · Feb 15, 2023 · 57504d9 · 57504d9
1 parent deda87b
commit 57504d9
Show file tree

Hide file tree

Showing 7 changed files with 201 additions and 77 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -463,7 +463,7 @@ commands:
       - run:
           name: Run core / C++ tests
           environment:
-            LD_LIBRARY_PATH: "/home/circleci/project/bazel-project/external/libtorch_pre_cxx11_abi/lib/:/home/circleci/project/bazel-project/external/tensorrt/lib/:/usr/local/cuda/lib64/:$LD_LIBRARY_PATH"
+            LD_LIBRARY_PATH: "/opt/circleci/.pyenv/versions/3.9.4/lib/python3.9/site-packages/torch_tensorrt.libs:/home/circleci/project/bazel-project/external/libtorch_pre_cxx11_abi/lib/:/home/circleci/project/bazel-project/external/tensorrt/lib/:/usr/local/cuda/lib64/:$LD_LIBRARY_PATH"
           command: |
             set -e
             mv toolchains/ci_workspaces/WORKSPACE.<< parameters.platform >> WORKSPACE
@@ -719,7 +719,7 @@ jobs:
           at: /tmp/dist
       - run:
           name: "Install torch-tensorrt"
-          command: pip3 install /tmp/dist/x86_64-linux/*
+          command: pip3 install /tmp/dist/x86_64-linux/*cp39-cp39*.whl
       - dump-test-env
       - test-ts-core
 
@@ -747,7 +747,7 @@ jobs:
           torch-build-index: << parameters.torch-build-index >>
       - run:
           name: "Install torch-tensorrt"
-          command: pip3 install --pre /tmp/dist/x86_64-linux/*
+          command: pip3 install --pre /tmp/dist/x86_64-linux/*cp39-cp39*.whl
       - dump-test-env
       - test-ts-py-api
 
@@ -777,7 +777,7 @@ jobs:
       #     command: export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH
       - run:
           name: "Install torch-tensorrt"
-          command: pip3 install --pre /tmp/dist/x86_64-linux/*
+          command: pip3 install --pre /tmp/dist/x86_64-linux/*cp39-cp39*.whl
       # We install torch after torch-trt because pip automatically enforces the version constraint otherwise
       - dump-test-env
       - test-fx
@@ -791,9 +791,10 @@ jobs:
         type: string
       torch-build-index:
         type: string
+    parallelism: 5
     machine:
       image: ubuntu-2004-cuda-11.4:202110-01
-    resource_class: xlarge
+    resource_class: gpu.nvidia.small
     steps:
       - when:
           condition: << parameters.enabled >>
@@ -805,18 +806,58 @@ jobs:
                   cd ~/project/py/
                   docker build -t torch_tensorrt_release_env --build-arg trt_version=<< pipeline.parameters.trt-version-short >> -f ci/Dockerfile.ci .
             - run:
-                name: Build Python packages and pre-cxx11-abi tarball
+                name: Build Python packages
                 command: |
                   cd ~/project/py/
                   cp ~/project/toolchains/ci_workspaces/WORKSPACE.x86_64.release.rhel ~/project/WORKSPACE
-                  docker run -it --rm -v ~/project:/workspace/project torch_tensorrt_release_env /bin/bash /workspace/project/py/ci/build_whl.sh
+                  docker run -it --rm --gpus all -v ~/project:/workspace/project torch_tensorrt_release_env /bin/bash -c "source /workspace/project/py/ci/build_whl.sh && $(circleci tests split --total=5 ~/project/py/ci/build_manifest.txt)"
+            - run:
+                name: Collect packages
+                command: |
+                  mkdir -p /tmp/dist/x86_64-linux
+                  cp -r ~/project/py/wheelhouse/* /tmp/dist/x86_64-linux
+            - persist_to_workspace:
+                root: /tmp/dist
+                paths:
+                  - x86_64-linux
+            - store_artifacts:
+                path: /tmp/dist/x86_64-linux
+                destination: x86_64-linux
+      - unless:
+          condition: << parameters.enabled >>
+          steps:
+            - run:
+                name: Skipped packaging
+                command: echo -e "Packaging stage not enabled"
+
+  # TODO: Merge this with above
+  package-x86_64-linux-cxx11-abi:
+    parameters:
+      enabled:
+        type: boolean
+        default: false
+      torch-build:
+        type: string
+      torch-build-index:
+        type: string
+    machine:
+      image: ubuntu-2004-cuda-11.4:202110-01
+    resource_class: xlarge
+    steps:
+      - when:
+          condition: << parameters.enabled >>
+          steps:
+            - checkout
             - create-env:
                 os: "ubuntu2004"
                 platform: "x86_64"
                 cudnn-version: << pipeline.parameters.cudnn-version >>
                 trt-version-short:  << pipeline.parameters.trt-version-short >>
                 bazel-version: << pipeline.parameters.bazel-version >>
                 bazel-platform: "x86_64"
+            - install-torch-from-index:
+                torch-build:  << parameters.torch-build >>
+                torch-build-index: << parameters.torch-build-index >>
             - run:
                 name: Build cxx11-abi tarball
                 command: |
@@ -825,13 +866,16 @@ jobs:
                   cp ~/project/toolchains/ci_workspaces/WORKSPACE.x86_64.release.ubuntu ~/project/WORKSPACE
                   bazel build //:libtorchtrt -c opt --noshow_progress
                   sudo chown -R $(whoami) ~/project/py
-                  CUDA_VERSION=$(cd ~/project/py/torch_tensorrt && python3 -c "from _version import __cuda_version__;print(__cuda_version__)")
-                  TORCHTRT_VERSION=$(cd ~/project/py/torch_tensorrt && python3 -c "from _version import __version__;print(__version__)")
-                  TRT_VERSION=$(cd ~/project/py/torch_tensorrt && python3 -c "from _version import __tensorrt_version__;print(__tensorrt_version__)")
-                  CUDNN_VERSION=$(cd ~/project/py/torch_tensorrt && python3 -c "from _version import __cudnn_version__;print(__cudnn_version__)")
+                  cd ~/project/py
+                  CUDA_VERSION=$(python3 -c "from versions import __cuda_version__;print(__cuda_version__)")
+                  TORCHTRT_VERSION=$(python3 -c "from versions import __version__;print(__version__)")
+                  TRT_VERSION=$(python3 -c "from versions import __tensorrt_version__;print(__tensorrt_version__)")
+                  CUDNN_VERSION=$(python3 -c "from versions import __cudnn_version__;print(__cudnn_version__)")
+                  TORCH_VERSION=$(python3 -c "from torch import __version__;print(__version__.split('+')[0])")
                   pip3 install --upgrade pip
                   pip3 install -r ~/project/py/requirements.txt
                   TORCH_VERSION=$(python3 -c "from torch import __version__;print(__version__.split('+')[0])")
+                  mkdir -p ~/project/py/dist/
                   cp ~/project/bazel-bin/libtorchtrt.tar.gz ~/project/py/dist/libtorchtrt-${TORCHTRT_VERSION}-cudnn${CUDNN_VERSION}-tensorrt${TRT_VERSION}-cuda${CUDA_VERSION}-libtorch${TORCH_VERSION}-x86_64-linux.tar.gz
             - run:
                 name: Collect packages
@@ -1113,7 +1157,13 @@ workflows:
     when: << pipeline.parameters.enable-packaging >>
     jobs:
 
-      - build-x86_64-linux:
+      - package-x86_64-linux:
+          enabled: << pipeline.parameters.enable-packaging >>
+          torch-build: << pipeline.parameters.torch-build >>
+          torch-build-index: << pipeline.parameters.torch-build-index >>
+
+      - package-x86_64-linux-cxx11-abi:
+          enabled: << pipeline.parameters.enable-packaging >>
           torch-build: << pipeline.parameters.torch-build >>
           torch-build-index: << pipeline.parameters.torch-build-index >>
 
@@ -1124,31 +1174,21 @@ workflows:
           trt-version-long: << pipeline.parameters.trt-version-long >>
           cudnn-version: << pipeline.parameters.cudnn-version >>
           requires:
-            - build-x86_64-linux
+            - package-x86_64-linux
 
       - test-py-ts-x86_64-linux:
           torch-build: << pipeline.parameters.torch-build >>
           torch-build-index: << pipeline.parameters.torch-build-index >>
           trt-version-long: << pipeline.parameters.trt-version-long >>
           requires:
-            - build-x86_64-linux
+            - package-x86_64-linux
 
       - test-py-fx-x86_64-linux:
           torch-build: << pipeline.parameters.torch-build >>
           torch-build-index: << pipeline.parameters.torch-build-index >>
           trt-version-long: << pipeline.parameters.trt-version-long >>
           requires:
-            - build-x86_64-linux
-
-
-      - package-x86_64-linux:
-          enabled: << pipeline.parameters.enable-packaging >>
-          torch-build: << pipeline.parameters.torch-build >>
-          torch-build-index: << pipeline.parameters.torch-build-index >>
-          #requires:
-            #- test-core-cpp-x86_64-linux
-            #- test-py-ts-x86_64-linux
-            #- test-py-fx-x86_64-linux
+            - package-x86_64-linux
 
   on-push:
     jobs:

diff --git a/py/ci/build_manifest.txt b/py/ci/build_manifest.txt
@@ -0,0 +1,5 @@
+py37
+py38
+py39
+py310
+libtorchtrt_pre_cxx11_abi
diff --git a/py/ci/build_whl.sh b/py/ci/build_whl.sh
@@ -2,36 +2,66 @@
 
 # Example usage: docker run -it -v$(pwd)/..:/workspace/TRTorch build_trtorch_wheel /bin/bash /workspace/TRTorch/py/build_whl.sh
 
-cd /workspace/project/py
-
 export CXX=g++
-export CUDA_HOME=/usr/local/cuda-11.6
-PROJECT_DIR=/workspace/project
+export CUDA_HOME=/usr/local/cuda-11.7
+export PROJECT_DIR=/workspace/project
 
 cp -r $CUDA_HOME /usr/local/cuda
 
-build_py37() {
-    /opt/python/cp37-cp37m/bin/python -m pip install -r requirements.txt
-    /opt/python/cp37-cp37m/bin/python setup.py bdist_wheel --release --ci
-    #auditwheel repair --plat manylinux2014_x86_64
+py37() {
+    cd /workspace/project/py
+    PY_BUILD_CODE=cp37-cp37m
+    PY_VERSION=3.7
+    PY_NAME=python${PY_VERSION}
+    PY_DIR=/opt/python/${PY_BUILD_CODE}
+    PY_PKG_DIR=${PY_DIR}/lib/${PY_NAME}/site-packages/
+    ${PY_DIR}/bin/python -m pip install --upgrade pip
+    ${PY_DIR}/bin/python -m pip install -r requirements.txt
+    ${PY_DIR}/bin/python -m pip install setuptools wheel auditwheel
+    ${PY_DIR}/bin/python setup.py bdist_wheel --release --ci
+    LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${PY_PKG_DIR}/torch/lib:${PY_PKG_DIR}/tensorrt/:${CUDA_HOME}/lib64:${CUDA_HOME}/lib64/stubs ${PY_DIR}/bin/python -m auditwheel repair $(cat ${PROJECT_DIR}/py/ci/soname_excludes.params) --plat manylinux_2_17_x86_64 dist/torch_tensorrt-*-${PY_BUILD_CODE}-linux_x86_64.whl
 }
 
-build_py38() {
-    /opt/python/cp38-cp38/bin/python -m pip install -r requirements.txt
-    /opt/python/cp38-cp38/bin/python setup.py bdist_wheel --release --ci
-    #auditwheel repair --plat manylinux2014_x86_64
+py38() {
+    cd /workspace/project/py
+    PY_BUILD_CODE=cp38-cp38
+    PY_VERSION=3.8
+    PY_NAME=python${PY_VERSION}
+    PY_DIR=/opt/python/${PY_BUILD_CODE}
+    PY_PKG_DIR=${PY_DIR}/lib/${PY_NAME}/site-packages/
+    ${PY_DIR}/bin/python -m pip install --upgrade pip
+    ${PY_DIR}/bin/python -m pip install -r requirements.txt
+    ${PY_DIR}/bin/python -m pip install setuptools wheel auditwheel
+    ${PY_DIR}/bin/python setup.py bdist_wheel --release --ci
+    LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${PY_PKG_DIR}/torch/lib:${PY_PKG_DIR}/tensorrt/:${CUDA_HOME}/lib64:${CUDA_HOME}/lib64/stubs ${PY_DIR}/bin/python -m auditwheel repair  $(cat ${PROJECT_DIR}/py/ci/soname_excludes.params) --plat manylinux_2_17_x86_64 dist/torch_tensorrt-*-${PY_BUILD_CODE}-linux_x86_64.whl
 }
 
-build_py39() {
-    /opt/python/cp39-cp39/bin/python -m pip install -r requirements.txt
-    /opt/python/cp39-cp39/bin/python setup.py bdist_wheel --release --ci
-    #auditwheel repair --plat manylinux2014_x86_64
+py39() {
+    cd /workspace/project/py
+    PY_BUILD_CODE=cp39-cp39
+    PY_VERSION=3.9
+    PY_NAME=python${PY_VERSION}
+    PY_DIR=/opt/python/${PY_BUILD_CODE}
+    PY_PKG_DIR=${PY_DIR}/lib/${PY_NAME}/site-packages/
+    ${PY_DIR}/bin/python -m pip install --upgrade pip
+    ${PY_DIR}/bin/python -m pip install -r requirements.txt
+    ${PY_DIR}/bin/python -m pip install setuptools wheel auditwheel
+    ${PY_DIR}/bin/python setup.py bdist_wheel --release --ci
+    LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${PY_PKG_DIR}/torch/lib:${PY_PKG_DIR}/tensorrt/:${CUDA_HOME}/lib64:${CUDA_HOME}/lib64/stubs ${PY_DIR}/bin/python -m auditwheel repair  $(cat ${PROJECT_DIR}/py/ci/soname_excludes.params) --plat manylinux_2_17_x86_64 dist/torch_tensorrt-*-${PY_BUILD_CODE}-linux_x86_64.whl
 }
 
-build_py310() {
-    /opt/python/cp310-cp310/bin/python -m pip install -r requirements.txt
-    /opt/python/cp310-cp310/bin/python setup.py bdist_wheel --release --ci
-    #auditwheel repair --plat manylinux2014_x86_64
+py310() {
+    cd /workspace/project/py
+    PY_BUILD_CODE=cp310-cp310
+    PY_VERSION=3.10
+    PY_NAME=python${PY_VERSION}
+    PY_DIR=/opt/python/${PY_BUILD_CODE}
+    PY_PKG_DIR=${PY_DIR}/lib/${PY_NAME}/site-packages/
+    ${PY_DIR}/bin/python -m pip install --upgrade pip
+    ${PY_DIR}/bin/python -m pip install -r requirements.txt
+    ${PY_DIR}/bin/python -m pip install setuptools wheel auditwheel
+    ${PY_DIR}/bin/python setup.py bdist_wheel --release --ci
+    LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${PY_PKG_DIR}/torch/lib:${PY_PKG_DIR}/tensorrt/:${CUDA_HOME}/lib64:${CUDA_HOME}/lib64/stubs ${PY_DIR}/bin/python -m auditwheel repair  $(cat ${PROJECT_DIR}/py/ci/soname_excludes.params) --plat manylinux_2_17_x86_64 dist/torch_tensorrt-*-${PY_BUILD_CODE}-linux_x86_64.whl
 }
 
 #build_py311() {
@@ -40,31 +70,42 @@ build_py310() {
 #    #auditwheel repair --plat manylinux2014_x86_64
 #}
 
-build_libtorchtrt() {
-    bazel clean
+libtorchtrt() {
+    cd /workspace/project/py
+    mkdir -p /workspace/project/py/wheelhouse
+    PY_BUILD_CODE=cp310-cp310
+    PY_VERSION=3.10
+    PY_NAME=python${PY_VERSION}
+    PY_DIR=/opt/python/${PY_BUILD_CODE}
+    PY_PKG_DIR=${PY_DIR}/lib/${PY_NAME}/site-packages/
+    ${PY_DIR}/bin/python -m pip install --upgrade pip
+    ${PY_DIR}/bin/python -m pip install -r requirements.txt
+    ${PY_DIR}/bin/python -m pip install setuptools wheel auditwheel
     bazel build //:libtorchtrt --platforms //toolchains:ci_rhel_x86_64_linux -c opt --noshow_progress
-    CUDA_VERSION=$(cd ${PROJECT_DIR}/py/torch_tensorrt && python3 -c "from _version import __cuda_version__;print(__cuda_version__)")
-    TORCHTRT_VERSION=$(cd ${PROJECT_DIR}/py/torch_tensorrt && python3 -c "from _version import __version__;print(__version__)")
-    TRT_VERSION=$(cd ${PROJECT_DIR}/py/torch_tensorrt && python3 -c "from _version import __tensorrt_version__;print(__tensorrt_version__)")
-    CUDNN_VERSION=$(cd ${PROJECT_DIR}/py/torch_tensorrt && python3 -c "from _version import __cudnn_version__;print(__cudnn_version__)")
-    TORCH_VERSION=$(/opt/python/cp310-cp310/bin/python -c "from torch import __version__;print(__version__.split('+')[0])")
-    cp ${PROJECT_DIR}/bazel-bin/libtorchtrt.tar.gz ${PROJECT_DIR}/py/dist/libtorchtrt-${TORCHTRT_VERSION}-cudnn${CUDNN_VERSION}-tensorrt${TRT_VERSION}-cuda${CUDA_VERSION}-libtorch${TORCH_VERSION}-x86_64-linux.tar.gz
+    CUDA_VERSION=$(cd ${PROJECT_DIR}/py && ${PY_DIR}/bin/python3 -c "from versions import __cuda_version__;print(__cuda_version__)")
+    TORCHTRT_VERSION=$(cd ${PROJECT_DIR}/py && ${PY_DIR}/bin/python3 -c "from versions import __version__;print(__version__)")
+    TRT_VERSION=$(cd ${PROJECT_DIR}/py && ${PY_DIR}/bin/python3 -c "from versions import __tensorrt_version__;print(__tensorrt_version__)")
+    CUDNN_VERSION=$(cd ${PROJECT_DIR}/py && ${PY_DIR}/bin/python3 -c "from versions import __cudnn_version__;print(__cudnn_version__)")
+    TORCH_VERSION=$(${PY_DIR}/bin/python -c "from torch import __version__;print(__version__.split('+')[0])")
+    cp ${PROJECT_DIR}/bazel-bin/libtorchtrt.tar.gz ${PROJECT_DIR}/py/wheelhouse/libtorchtrt-${TORCHTRT_VERSION}-cudnn${CUDNN_VERSION}-tensorrt${TRT_VERSION}-cuda${CUDA_VERSION}-libtorch${TORCH_VERSION}-x86_64-linux.tar.gz
 }
 
-build_libtorchtrt_pre_cxx11_abi() {
+libtorchtrt_pre_cxx11_abi() {
+    cd /workspace/project/py
+    mkdir -p /workspace/project/py/wheelhouse
+    PY_BUILD_CODE=cp310-cp310
+    PY_VERSION=3.10
+    PY_NAME=python${PY_VERSION}
+    PY_DIR=/opt/python/${PY_BUILD_CODE}
+    PY_PKG_DIR=${PY_DIR}/lib/${PY_NAME}/site-packages/
+    ${PY_DIR}/bin/python -m pip install --upgrade pip
+    ${PY_DIR}/bin/python -m pip install -r requirements.txt
+    ${PY_DIR}/bin/python -m pip install setuptools wheel auditwheel
     bazel build //:libtorchtrt --config pre_cxx11_abi --platforms //toolchains:ci_rhel_x86_64_linux -c opt --noshow_progress
-    CUDA_VERSION=$(cd ${PROJECT_DIR}/py/torch_tensorrt && python3 -c "from _version import __cuda_version__;print(__cuda_version__)")
-    TORCHTRT_VERSION=$(cd ${PROJECT_DIR}/py/torch_tensorrt && python3 -c "from _version import __version__;print(__version__)")
-    TRT_VERSION=$(cd ${PROJECT_DIR}/py/torch_tensorrt && python3 -c "from _version import __tensorrt_version__;print(__tensorrt_version__)")
-    CUDNN_VERSION=$(cd ${PROJECT_DIR}/py/torch_tensorrt && python3 -c "from _version import __cudnn_version__;print(__cudnn_version__)")
-    TORCH_VERSION=$(/opt/python/cp310-cp310/bin/python -c "from torch import __version__;print(__version__.split('+')[0])")
-    cp ${PROJECT_DIR}/bazel-bin/libtorchtrt.tar.gz ${PROJECT_DIR}/py/dist/libtorchtrt-${TORCHTRT_VERSION}-pre-cxx11-abi-cudnn${CUDNN_VERSION}-tensorrt${TRT_VERSION}-cuda${CUDA_VERSION}-libtorch${TORCH_VERSION}-x86_64-linux.tar.gz
-}
-
-build_py37
-build_py38
-build_py39
-build_py310
-#build_py311
-build_libtorchtrt_pre_cxx11_abi
-#build_libtorchtrt
+    CUDA_VERSION=$(cd ${PROJECT_DIR}/py && ${PY_DIR}/bin/python3 -c "from versions import __cuda_version__;print(__cuda_version__)")
+    TORCHTRT_VERSION=$(cd ${PROJECT_DIR}/py && ${PY_DIR}/bin/python3 -c "from versions import __version__;print(__version__)")
+    TRT_VERSION=$(cd ${PROJECT_DIR}/py && ${PY_DIR}/bin/python3 -c "from versions import __tensorrt_version__;print(__tensorrt_version__)")
+    CUDNN_VERSION=$(cd ${PROJECT_DIR}/py && ${PY_DIR}/bin/python3 -c "from versions import __cudnn_version__;print(__cudnn_version__)")
+    TORCH_VERSION=$(${PY_DIR}/bin/python -c "from torch import __version__;print(__version__.split('+')[0])")
+    cp ${PROJECT_DIR}/bazel-bin/libtorchtrt.tar.gz ${PROJECT_DIR}/py/wheelhouse/libtorchtrt-${TORCHTRT_VERSION}-pre-cxx11-abi-cudnn${CUDNN_VERSION}-tensorrt${TRT_VERSION}-cuda${CUDA_VERSION}-libtorch${TORCH_VERSION}-x86_64-linux.tar.gz
+}
diff --git a/py/ci/soname_excludes.params b/py/ci/soname_excludes.params
@@ -0,0 +1,33 @@
+--exclude libc10_cuda.so
+--exclude libc10.so
+--exclude libcaffe2_nvrtc.so
+--exclude libcublasLt.so
+--exclude libcuda.so
+--exclude libnvrtc.so
+--exclude libnvToolsExt.so
+--exclude libshm.so
+--exclude libtorch_cpu.so
+--exclude libtorch_cuda_cpp.so
+--exclude libtorch_cuda_cu.so
+--exclude libtorch_cuda.so
+--exclude libtorch_global_deps.so
+--exclude libtorch_python.so
+--exclude libtorch.so
+--exclude libcudart.so
+
+--exclude libnvToolsExt-847d78f2.so.1
+--exclude libnvToolsExt-847d78f2.so
+--exclude libcudart-e409450e.so.11.0
+--exclude libcudart-e409450e.so
+--exclude libgomp-a34b3233.so.1
+--exclude libgomp-a34b3233.so
+--exclude libcudart.so.11
+--exclude libcudart.so.11.7.60
+--exclude libnvrtc.so.11.2
+--exclude libnvinfer_plugin.so.8
+--exclude libcublas.so.11
+--exclude libcuda.so.1
+--exclude libcuda.so.515
+--exclude libcublasLt.so.11
+--exclude libnvinfer.so.8
+--exclude libcudnn.so.8