fix: Update Docker build to automatically adapt Torch version (#1732)

Author: gs-olive

docker/Dockerfile

@@ -2,6 +2,9 @@
 ARG BASE_IMG=nvidia/cuda:11.7.1-devel-ubuntu20.04
 FROM ${BASE_IMG} as base
 
+ARG USE_CXX11_ABI
+ENV USE_CXX11=${USE_CXX11_ABI}
+
 # Install basic dependencies
 RUN apt-get update
 RUN DEBIAN_FRONTEND=noninteractive apt install -y build-essential manpages-dev wget zlib1g software-properties-common git
@@ -12,9 +15,6 @@ RUN ln -s /usr/bin/python3.8 /usr/bin/python
 RUN python get-pip.py
 RUN pip3 install wheel
 
-# Install Pytorch
-RUN pip3 install torch==2.0.0.dev20230103+cu117 torchvision==0.15.0.dev20230103+cu117 --extra-index-url https://download.pytorch.org/whl/nightly/cu117
-
 # Install CUDNN + TensorRT
 RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
 RUN mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
@@ -32,8 +32,7 @@ RUN apt-get update
 RUN apt-get install -y libnvinfer8=8.5.1* libnvinfer-plugin8=8.5.1* libnvinfer-dev=8.5.1* libnvinfer-plugin-dev=8.5.1* libnvonnxparsers8=8.5.1-1* libnvonnxparsers-dev=8.5.1-1* libnvparsers8=8.5.1-1*  libnvparsers-dev=8.5.1-1*
 
 # Setup Bazel
-ARG BAZEL_VERSION=5.2.0
-RUN wget -q https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION}/bazel-${BAZEL_VERSION}-linux-x86_64 -O /usr/bin/bazel \
+RUN wget -q https://github.com/bazelbuild/bazelisk/releases/download/v1.16.0/bazelisk-linux-amd64 -O /usr/bin/bazel \
  && chmod a+x /usr/bin/bazel
 
 # Build Torch-TensorRT in an auxillary container
@@ -44,7 +43,6 @@ ARG TARGETARCH="amd64"
 
 RUN apt-get install -y python3-setuptools
 RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub
-RUN  apt-get update
 
 RUN apt-get update && apt-get install -y --no-install-recommends locales ninja-build && rm -rf /var/lib/apt/lists/* && locale-gen en_US.UTF-8
 
@@ -64,11 +62,9 @@ COPY . /opt/torch_tensorrt
 COPY --from=torch-tensorrt-builder  /workspace/torch_tensorrt/src/py/dist/ .
 
 RUN cp /opt/torch_tensorrt/docker/WORKSPACE.docker /opt/torch_tensorrt/WORKSPACE
+RUN pip install -r /opt/torch_tensorrt/py/requirements.txt
 RUN pip3 install *.whl && rm -fr /workspace/torch_tensorrt/py/dist/* *.whl
 
-# Install native tensorrt python package required by torch_tensorrt whl file
-RUN pip install tensorrt==8.5.1.7
-
 WORKDIR /opt/torch_tensorrt
 ENV LD_LIBRARY_PATH /usr/local/lib/python3.8/dist-packages/torch/lib:/usr/local/lib/python3.8/dist-packages/torch_tensorrt/lib:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH}
 ENV PATH /usr/local/lib/python3.8/dist-packages/torch_tensorrt/bin:${PATH}

docker/README.md

@@ -5,7 +5,8 @@
 * `Dockerfile` currently uses the exact library versions (Torch, CUDA, CUDNN, TensorRT) listed in <a href="https://github.com/pytorch/TensorRT#dependencies">dependencies</a> to build Torch-TensorRT.
 
 * This `Dockerfile` installs `pre-cxx11-abi` versions of Pytorch and builds Torch-TRT using `pre-cxx11-abi` libtorch as well.
-Note: To install `cxx11_abi` version of Torch-TensorRT, enable `USE_CXX11=1` flag so that `dist-build.sh` can build it accordingly.
+
+Note: By default the container uses the `pre-cxx11-abi` version of Torch + Torch-TRT. If you are using a workflow that requires a build of PyTorch on the CXX11 ABI (e.g. using the PyTorch NGC containers as a base image), add the Docker build argument: `--build-arg USE_CXX11_ABI=1`
 
 ### Dependencies
 
@@ -38,4 +39,4 @@ bazel test //tests/core/conversion/converters:test_activation --compilation_mode
 
 ### Pytorch NGC containers
 
-We also ship Torch-TensorRT in <a href="https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch">Pytorch NGC containers </a>. Release notes for these containers can be found <a href="https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/index.html">here</a>. Check out `release/ngc/23.XX` branch of Torch-TensorRT for source code that gets shipped with `23.XX` version of Pytorch NGC container.
+We also ship Torch-TensorRT in <a href="https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch">Pytorch NGC containers </a>. Release notes for these containers can be found <a href="https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/index.html">here</a>. Check out `release/ngc/23.XX` branch of Torch-TensorRT for source code that gets shipped with `23.XX` version of Pytorch NGC container.

docker/WORKSPACE.docker

@@ -48,20 +48,16 @@ new_local_repository(
 # Tarballs and fetched dependencies (default - use in cases when building from precompiled bin and tarballs)
 #############################################################################################################
 
-http_archive(
+new_local_repository(
     name = "libtorch",
-    build_file = "@//third_party/libtorch:BUILD",
-    sha256 = "59b8b5e1954a86d50b79c13f06398d385b200da13e37a08ecf31d3c62e5ca127",
-    strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/nightly/cu117/libtorch-cxx11-abi-shared-with-deps-2.0.0.dev20230103%2Bcu117.zip"],
+    path = "/usr/local/lib/python3.8/dist-packages/torch/",
+    build_file = "third_party/libtorch/BUILD"
 )
 
-http_archive(
+new_local_repository(
     name = "libtorch_pre_cxx11_abi",
-    build_file = "@//third_party/libtorch:BUILD",
-    sha256 = "e260fc7476be89d1650953e8643e9f7363845f5a52de4bab87ac0e619c1f6ad4",
-    strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/nightly/cu117/libtorch-shared-with-deps-2.0.0.dev20230103%2Bcu117.zip"],
+    path = "/usr/local/lib/python3.8/dist-packages/torch/",
+    build_file = "third_party/libtorch/BUILD"
 )
 
 ####################################################################################

py/requirements.txt

@@ -1,4 +1,5 @@
 numpy
+packaging
 pybind11==2.6.2
 --extra-index-url https://download.pytorch.org/whl/nightly/cu117
 torch==2.1.0.dev20230314+cu117