feat: Dockerfile restructure for 1.3 release (#1492)

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

README.md

@@ -110,7 +110,7 @@ torch.jit.save(trt_ts_module, "trt_torchscript_module.ts") # save the TRT embedd
 
 ### Dependencies
 
-These are the following dependencies used to verify the testcases. Torch-TensorRT can work with other versions, but the tests are not guaranteed to pass.
+These are the following dependencies used to verify the testcases. Check out `py/requirements.txt` for python dependencies. Torch-TensorRT can work with other versions, but the tests are not guaranteed to pass.
 
 - Bazel 5.2.0
 - Libtorch 1.13.0 (built with CUDA 11.7)

docker/Dockerfile

@@ -1,62 +1,80 @@
-ARG BASE=22.04
-ARG BASE_IMG=nvcr.io/nvidia/tensorrt:${BASE}-py3
+# Base image starts with CUDA
+ARG BASE_IMG=nvidia/cuda:11.7.1-devel-ubuntu18.04
 FROM ${BASE_IMG} as base
 
-FROM base as torch-tensorrt-builder-base
+# Install basic dependencies
+RUN apt-get update
+RUN apt install -y build-essential manpages-dev wget zlib1g software-properties-common git
+RUN add-apt-repository ppa:deadsnakes/ppa
+RUN apt install -y python3.8 python3.8-distutils python3.8-dev
+RUN wget https://bootstrap.pypa.io/get-pip.py
+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==1.13.0 torchvision==0.14.0
+
+# Install CUDNN + TensorRT
+RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-ubuntu1804.pin
+RUN mv cuda-ubuntu1804.pin /etc/apt/preferences.d/cuda-repository-pin-600
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
+RUN apt-key adv --keyserver keyserver.ubuntu.com --recv-keys 536F8F1DE80F6A35
+RUN apt-key adv --keyserver keyserver.ubuntu.com --recv-keys A4B469963BF863CC
+RUN add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/ /"
+RUN apt-get update
+RUN apt-get install -y libcudnn8=8.5.0* libcudnn8-dev=8.5.0*
+
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub
+RUN add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/ /"
+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 \
+ && chmod a+x /usr/bin/bazel
+
+# Remove cuda symlink to avoid bazel circle symlink errors
+RUN rm /usr/local/cuda-11.7/cuda-11.7
 
-# Removing any bazel or torch-tensorrt pre-installed from the base image
-RUN rm -rf /opt/pytorch/torch_tensorrt /usr/bin/bazel
+# Build Torch-TensorRT in an auxillary container
+FROM base as torch-tensorrt-builder-base
 
 ARG ARCH="x86_64"
 ARG TARGETARCH="amd64"
-ARG BAZEL_VERSION=5.2.0
-
-RUN [[ "$TARGETARCH" == "amd64" ]] && ARCH="x86_64" || ARCH="${TARGETARCH}" \
- && wget -q https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION}/bazel-${BAZEL_VERSION}-linux-${ARCH} -O /usr/bin/bazel \
- && chmod a+x /usr/bin/bazel
 
-# Workaround for bazel expecting both static and shared versions, we only use shared libraries inside container
-RUN touch /usr/lib/$HOSTTYPE-linux-gnu/libnvinfer_static.a
-
-RUN rm -rf /usr/local/cuda/lib* /usr/local/cuda/include \
-  && ln -sf /usr/local/cuda/targets/$HOSTTYPE-linux/lib /usr/local/cuda/lib64 \
-  && ln -sf /usr/local/cuda/targets/$HOSTTYPE-linux/include /usr/local/cuda/include
+RUN apt-get install -y python3-setuptools
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/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
 
 FROM torch-tensorrt-builder-base as torch-tensorrt-builder
 
-# Removing any bazel or torch-tensorrt pre-installed from the base image
-RUN rm -rf /opt/pytorch/torch_tensorrt
-
 COPY . /workspace/torch_tensorrt/src
 WORKDIR /workspace/torch_tensorrt/src
 RUN cp ./docker/WORKSPACE.docker WORKSPACE
 
 # This script builds both libtorchtrt bin/lib/include tarball and the Python wheel, in dist/
 RUN ./docker/dist-build.sh
 
+# Copy and install Torch-TRT into the main container
 FROM base as torch-tensorrt
 
-# Removing any bazel or torch-tensorrt pre-installed from the base image
-RUN rm -rf /opt/pytorch/torch_tensorrt
-
-# copy source repo
-COPY . /workspace/torch_tensorrt
+COPY . /opt/torch_tensorrt
 COPY --from=torch-tensorrt-builder  /workspace/torch_tensorrt/src/py/dist/ .
 
-RUN pip3 install ipywidgets --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org
-RUN jupyter nbextension enable --py widgetsnbextension
-
+RUN cp /opt/torch_tensorrt/docker/WORKSPACE.docker /opt/torch_tensorrt/WORKSPACE
 RUN pip3 install *.whl && rm -fr /workspace/torch_tensorrt/py/dist/* *.whl
 
-ENV LD_LIBRARY_PATH /usr/local/lib/python3.8/dist-packages/torch/lib:/usr/local/lib/python3.8/dist-packages/torch_tensorrt/lib:${LD_LIBRARY_PATH}
+# Install native tensorrt python package required by torch_tensorrt whl file
+RUN pip install nvidia-pyindex
+RUN pip install nvidia-tensorrt
+
+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}
-#
-WORKDIR /workspace
-RUN mv /workspace/torch_tensorrt /opt/torch_tensorrt
-RUN cp /opt/torch_tensorrt/docker/WORKSPACE.docker /opt/torch_tensorrt/WORKSPACE
-RUN mkdir torch_tensorrt
-RUN ln -s /opt/torch_tensorrt/notebooks /workspace/torch_tensorrt/notebooks
 
 CMD /bin/bash

docker/Dockerfile.ngc

@@ -1,19 +1,41 @@
 # Building a Torch-TensorRT container
 
-### Install Docker and NVIDIA Container Toolkit
+* Use `Dockerfile` to build a container which provides the exact development environment that our master branch is usually tested against.
 
-https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html
+* `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.
 
-### Build Container
+* 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.
+
+### Dependencies
+
+* Install nvidia-docker by following https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#docker
+
+### Instructions
 
 > From root of Torch-TensorRT repo
 
+Build:
+```
+DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile -t torch_tensorrt:latest .
 ```
-# Build:
-DOCKER_BUILDKIT=1 docker build --build-arg BASE={TensorRT Base Container Version} -f docker/Dockerfile -t torch_tensorrt1.0:latest .
 
-# Run:
-docker run --gpus all -it \
-	--shm-size=8gb --env="DISPLAY" --volume="/tmp/.X11-unix:/tmp/.X11-unix:rw" \
-	--name=torch_tensorrt1.0 --ipc=host --net=host torch_tensorrt1.0:latest
+Run:
 ```
+nvidia-docker run --gpus all -it --shm-size=8gb --env="DISPLAY" --volume="/tmp/.X11-unix:/tmp/.X11-unix:rw" --name=torch_tensorrt --ipc=host --net=host torch_tensorrt:latest
+```
+
+Test:
+
+
+You can run any converter test to verify if Torch-TRT built sucessfully inside the container. Once you launch the container, you can run
+```
+bazel test //tests/core/conversion/converters:test_activation --compilation_mode=opt --test_output=summary --config use_precompiled_torchtrt --config pre_cxx11_abi
+```
+
+* `--config use_precompiled_torchtrt` : Indicates bazel to use pre-installed Torch-TRT library to test an application.
+* `--config pre_cxx11_abi` : This flag ensures `bazel test` uses `pre_cxx11_abi` version of `libtorch`. Use this flag corresponding to the ABI format of your Torch-TensorRT installation.
+
+### 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/22.XX` branch of Torch-TensorRT for source code that gets shipped with `22.XX` version of Pytorch NGC container.

docker/README.md

@@ -31,11 +31,17 @@ git_repository(
     shallow_since = "1570114335 -0400",
 )
 
+# External dependency for torch_tensorrt if you already have precompiled binaries.
+local_repository(
+    name = "torch_tensorrt",
+    path = "/usr/local/lib/python3.8/dist-packages/torch_tensorrt",
+)
+
 # CUDA should be installed on the system locally
 new_local_repository(
     name = "cuda",
-    path = "/usr/local/cuda/",
     build_file = "@//third_party/cuda:BUILD",
+    path = "/usr/local/cuda-11.7/",
 )
 
 new_local_repository(
@@ -64,6 +70,7 @@ http_archive(
     urls = ["https://download.pytorch.org/libtorch/cu117/libtorch-shared-with-deps-1.13.0%2Bcu117.zip"],
 )
 
+
 ####################################################################################
 # Locally installed dependencies (use in cases of custom dependencies or aarch64)
 ####################################################################################

docker/WORKSPACE.docker

@@ -3,9 +3,9 @@
 TOP_DIR=$(cd $(dirname $0); pwd)/..
 
 if [[ -z "${USE_CXX11}" ]]; then
-    BUILD_CMD="python3 setup.py bdist_wheel"
+    BUILD_CMD="python setup.py bdist_wheel"
 else
-    BUILD_CMD="python3 setup.py bdist_wheel  --use-cxx11-abi"
+    BUILD_CMD="python setup.py bdist_wheel  --use-cxx11-abi"
 fi
 
 cd ${TOP_DIR} \