From 34fd7fe8eb58787318b463a46ebe889f89145d8c Mon Sep 17 00:00:00 2001
From: gs-olive <113141689+gs-olive@users.noreply.github.com>
Date: Mon, 3 Apr 2023 17:58:37 -0700
Subject: [PATCH 1/4] feat: Upgrade Docker build to use custom TRT + CUDNN
- Upgrade Ubuntu version to 22.04
- Make custom build args for TensorRT and CUDNN versions,
user-specifiable in a.b.c version format
- Make build args required, without defaults, to improve code
cleanliness and reduce amount of changes needed when versions shift
- Provide clear error messages when the build args are not provided by
the user
- Add documentation of the change in the Docker README
---
docker/Dockerfile | 29 ++++++++++++++++++-----------
docker/README.md | 7 +++++--
docker/dist-build.sh | 4 ++--
3 files changed, 25 insertions(+), 15 deletions(-)
diff --git a/docker/Dockerfile b/docker/Dockerfile
index ab804ed6c5..0bee5fc49e 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,13 +1,19 @@
# Base image starts with CUDA
-ARG BASE_IMG=nvidia/cuda:11.7.1-devel-ubuntu20.04
+ARG BASE_IMG=nvidia/cuda:11.7.1-devel-ubuntu22.04
FROM ${BASE_IMG} as base
+ARG TENSORRT_VERSION
+RUN test -n "$TENSORRT_VERSION" || (echo "No tensorrt version specified, please use --build-arg TENSORRT_VERSION==x.y.z to specify a version." && exit 1)
+ARG CUDNN_VERSION
+RUN test -n "$CUDNN_VERSION" || (echo "No cudnn version specified, please use --build-arg CUDNN_VERSION==x.y.z to specify a version." && exit 1)
+
ARG USE_CXX11_ABI
ENV USE_CXX11=${USE_CXX11_ABI}
+ENV DEBIAN_FRONTEND=noninteractive
# Install basic dependencies
RUN apt-get update
-RUN DEBIAN_FRONTEND=noninteractive apt install -y build-essential manpages-dev wget zlib1g software-properties-common git
+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
@@ -16,20 +22,20 @@ RUN python get-pip.py
RUN pip3 install wheel
# 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
-RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub
+RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
+RUN mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/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/ubuntu2004/x86_64/ /"
+RUN add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /"
RUN apt-get update
-RUN apt-get install -y libcudnn8=8.5.0* libcudnn8-dev=8.5.0*
+RUN apt-get install -y libcudnn8=${CUDNN_VERSION}* libcudnn8-dev=${CUDNN_VERSION}*
-RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub
-RUN add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/ /"
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub
+RUN add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/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*
+RUN apt-get install -y libnvinfer8=${TENSORRT_VERSION}* libnvinfer-plugin8=${TENSORRT_VERSION}* libnvinfer-dev=${TENSORRT_VERSION}* libnvinfer-plugin-dev=${TENSORRT_VERSION}* libnvonnxparsers8=${TENSORRT_VERSION}-1* libnvonnxparsers-dev=${TENSORRT_VERSION}-1* libnvparsers8=${TENSORRT_VERSION}-1* libnvparsers-dev=${TENSORRT_VERSION}-1*
# Setup Bazel
RUN wget -q https://github.com/bazelbuild/bazelisk/releases/download/v1.16.0/bazelisk-linux-amd64 -O /usr/bin/bazel \
@@ -42,7 +48,7 @@ ARG ARCH="x86_64"
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-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub
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
@@ -63,6 +69,7 @@ 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 pip install tensorrt==${TENSORRT_VERSION}.*
RUN pip3 install *.whl && rm -fr /workspace/torch_tensorrt/py/dist/* *.whl
WORKDIR /opt/torch_tensorrt
diff --git a/docker/README.md b/docker/README.md
index 3f69e0d431..9c164a9d03 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -2,7 +2,7 @@
* Use `Dockerfile` to build a container which provides the exact development environment that our master branch is usually tested against.
-* `Dockerfile` currently uses the exact library versions (Torch, CUDA, CUDNN, TensorRT) listed in dependencies to build Torch-TensorRT.
+* The `Dockerfile` currently uses Bazelisk to select the Bazel version, and uses the exact library versions of Torch and CUDA listed in dependencies. The desired versions of CUDNN and TensorRT must be specified as build-args, with major, minor, and patch versions as in: `--build-arg TENSORRT_VERSION=a.b.c --build-arg CUDNN_VERSION=x.y.z`
* This `Dockerfile` installs `pre-cxx11-abi` versions of Pytorch and builds Torch-TRT using `pre-cxx11-abi` libtorch as well.
@@ -14,11 +14,14 @@ Note: By default the container uses the `pre-cxx11-abi` version of Torch + Torch
### Instructions
+- The example below uses CUDNN 8.5.0 and TensorRT 8.5.1
+- See dependencies for a list of current default dependencies.
+
> From root of Torch-TensorRT repo
Build:
```
-DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile -t torch_tensorrt:latest .
+DOCKER_BUILDKIT=1 docker build --build-arg TENSORRT_VERSION=8.5.1 --build-arg CUDNN_VERSION=8.5.0 -f docker/Dockerfile -t torch_tensorrt:latest .
```
Run:
diff --git a/docker/dist-build.sh b/docker/dist-build.sh
index cf209258be..1655cf2c61 100755
--- a/docker/dist-build.sh
+++ b/docker/dist-build.sh
@@ -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} \
From dfec229fd12a6e454d9b7ad0dfbe7fbb7a167bc3 Mon Sep 17 00:00:00 2001
From: George S <113141689+gs-olive@users.noreply.github.com>
Date: Mon, 3 Apr 2023 18:51:20 -0700
Subject: [PATCH 2/4] Minor fix: Update error message
---
docker/Dockerfile | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 0bee5fc49e..f449a8439a 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -3,9 +3,9 @@ ARG BASE_IMG=nvidia/cuda:11.7.1-devel-ubuntu22.04
FROM ${BASE_IMG} as base
ARG TENSORRT_VERSION
-RUN test -n "$TENSORRT_VERSION" || (echo "No tensorrt version specified, please use --build-arg TENSORRT_VERSION==x.y.z to specify a version." && exit 1)
+RUN test -n "$TENSORRT_VERSION" || (echo "No tensorrt version specified, please use --build-arg TENSORRT_VERSION=x.y.z to specify a version." && exit 1)
ARG CUDNN_VERSION
-RUN test -n "$CUDNN_VERSION" || (echo "No cudnn version specified, please use --build-arg CUDNN_VERSION==x.y.z to specify a version." && exit 1)
+RUN test -n "$CUDNN_VERSION" || (echo "No cudnn version specified, please use --build-arg CUDNN_VERSION=x.y.z to specify a version." && exit 1)
ARG USE_CXX11_ABI
ENV USE_CXX11=${USE_CXX11_ABI}
From a26438e50b46d12830bf721f5ca55cca0c662ea1 Mon Sep 17 00:00:00 2001
From: gs-olive <113141689+gs-olive@users.noreply.github.com>
Date: Thu, 6 Apr 2023 21:09:43 -0700
Subject: [PATCH 3/4] feat: Add option to specify Python distribution
- Add support for pyenv in Docker file, allowing specification of a
custom Python version
- Update WORKSPACE files to automatically synchronize regardless of
Python version choice, using symbolic links
- Update dependencies accordingly
- Add documentation for changing the base image of the Docker container
---
docker/Dockerfile | 50 +++++++++++++++++++++++++++--------------
docker/README.md | 4 +++-
docker/WORKSPACE.docker | 4 ++--
docker/WORKSPACE.ngc | 6 ++---
docker/dist-build.sh | 14 +++++++-----
5 files changed, 50 insertions(+), 28 deletions(-)
diff --git a/docker/Dockerfile b/docker/Dockerfile
index f449a8439a..1e86bfc182 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -7,21 +7,30 @@ RUN test -n "$TENSORRT_VERSION" || (echo "No tensorrt version specified, please
ARG CUDNN_VERSION
RUN test -n "$CUDNN_VERSION" || (echo "No cudnn version specified, please use --build-arg CUDNN_VERSION=x.y.z to specify a version." && exit 1)
+ARG PYTHON_VERSION=3.10
+ENV PYTHON_VERSION=${PYTHON_VERSION}
+
ARG USE_CXX11_ABI
ENV USE_CXX11=${USE_CXX11_ABI}
ENV DEBIAN_FRONTEND=noninteractive
# 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 CUDNN + TensorRT
+RUN apt install -y build-essential manpages-dev wget zlib1g software-properties-common git libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget ca-certificates curl llvm libncurses5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev mecab-ipadic-utf8
+
+# Install PyEnv and desired Python version
+ENV HOME="/root"
+ENV PYENV_DIR="$HOME/.pyenv"
+ENV PATH="$PYENV_DIR/shims:$PYENV_DIR/bin:$PATH"
+RUN wget -L https://github.com/pyenv/pyenv-installer/raw/master/bin/pyenv-installer &&\
+ chmod 755 pyenv-installer &&\
+ bash pyenv-installer &&\
+ eval "$(pyenv init -)"
+
+RUN pyenv install -v ${PYTHON_VERSION}
+RUN pyenv global ${PYTHON_VERSION}
+
+# Install CUDNN + TensorRT + dependencies
RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
RUN mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/7fa2af80.pub
@@ -37,9 +46,9 @@ RUN apt-get update
RUN apt-get install -y libnvinfer8=${TENSORRT_VERSION}* libnvinfer-plugin8=${TENSORRT_VERSION}* libnvinfer-dev=${TENSORRT_VERSION}* libnvinfer-plugin-dev=${TENSORRT_VERSION}* libnvonnxparsers8=${TENSORRT_VERSION}-1* libnvonnxparsers-dev=${TENSORRT_VERSION}-1* libnvparsers8=${TENSORRT_VERSION}-1* libnvparsers-dev=${TENSORRT_VERSION}-1*
-# Setup 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
+# Setup Bazel via Bazelisk
+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
FROM base as torch-tensorrt-builder-base
@@ -50,7 +59,10 @@ ARG TARGETARCH="amd64"
RUN apt-get install -y python3-setuptools
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub
-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
+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
@@ -58,8 +70,11 @@ COPY . /workspace/torch_tensorrt/src
WORKDIR /workspace/torch_tensorrt/src
RUN cp ./docker/WORKSPACE.docker WORKSPACE
+# Symlink the path pyenv is using for python with the /opt directory for package sourcing
+RUN ln -s "`pyenv which python | xargs dirname | xargs dirname`/lib/python$PYTHON_VERSION/site-packages" "/opt"
+
# This script builds both libtorchtrt bin/lib/include tarball and the Python wheel, in dist/
-RUN ./docker/dist-build.sh
+RUN bash ./docker/dist-build.sh
# Copy and install Torch-TRT into the main container
FROM base as torch-tensorrt
@@ -70,10 +85,11 @@ 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 pip install tensorrt==${TENSORRT_VERSION}.*
-RUN pip3 install *.whl && rm -fr /workspace/torch_tensorrt/py/dist/* *.whl
+RUN pip install *.whl && rm -fr /workspace/torch_tensorrt/py/dist/* *.whl
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}
+
+ENV LD_LIBRARY_PATH /opt/site-packages/torch/lib:/opt/site-packages/torch_tensorrt/lib:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH}
+ENV PATH /opt/site-packages/torch_tensorrt/bin:${PATH}
CMD /bin/bash
diff --git a/docker/README.md b/docker/README.md
index 9c164a9d03..a20a613553 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -2,7 +2,9 @@
* Use `Dockerfile` to build a container which provides the exact development environment that our master branch is usually tested against.
-* The `Dockerfile` currently uses Bazelisk to select the Bazel version, and uses the exact library versions of Torch and CUDA listed in dependencies. The desired versions of CUDNN and TensorRT must be specified as build-args, with major, minor, and patch versions as in: `--build-arg TENSORRT_VERSION=a.b.c --build-arg CUDNN_VERSION=x.y.z`
+* The `Dockerfile` currently uses Bazelisk to select the Bazel version, and uses the exact library versions of Torch and CUDA listed in dependencies.
+ * The desired versions of CUDNN and TensorRT must be specified as build-args, with major, minor, and patch versions as in: `--build-arg TENSORRT_VERSION=a.b.c --build-arg CUDNN_VERSION=x.y.z`
+ * The desired base image be changed by explicitly setting a base image, as in `--build-arg BASE_IMG=nvidia/cuda:11.7.1-devel-ubuntu22.04`, though this is optional
* This `Dockerfile` installs `pre-cxx11-abi` versions of Pytorch and builds Torch-TRT using `pre-cxx11-abi` libtorch as well.
diff --git a/docker/WORKSPACE.docker b/docker/WORKSPACE.docker
index d4a2f83d85..8a391c2075 100755
--- a/docker/WORKSPACE.docker
+++ b/docker/WORKSPACE.docker
@@ -50,13 +50,13 @@ new_local_repository(
new_local_repository(
name = "libtorch",
- path = "/usr/local/lib/python3.8/dist-packages/torch/",
+ path = "/opt/site-packages/torch/",
build_file = "third_party/libtorch/BUILD"
)
new_local_repository(
name = "libtorch_pre_cxx11_abi",
- path = "/usr/local/lib/python3.8/dist-packages/torch/",
+ path = "/opt/site-packages/torch/",
build_file = "third_party/libtorch/BUILD"
)
diff --git a/docker/WORKSPACE.ngc b/docker/WORKSPACE.ngc
index 7fd4aad4e6..79bed14dc9 100755
--- a/docker/WORKSPACE.ngc
+++ b/docker/WORKSPACE.ngc
@@ -33,7 +33,7 @@ git_repository(
# This is currently used in pytorch NGC container CI testing.
local_repository(
name = "torch_tensorrt",
- path = "/usr/local/lib/python3.8/dist-packages/torch_tensorrt"
+ path = "/usr/local/lib/python/site-packages/torch_tensorrt/"
)
# CUDA should be installed on the system locally
@@ -55,13 +55,13 @@ new_local_repository(
new_local_repository(
name = "libtorch",
- path = "/usr/local/lib/python3.8/dist-packages/torch",
+ path = "/opt/site-packages/torch/",
build_file = "third_party/libtorch/BUILD"
)
new_local_repository(
name = "libtorch_pre_cxx11_abi",
- path = "/usr/local/lib/python3.8/dist-packages/torch",
+ path = "/opt/site-packages/torch/",
build_file = "third_party/libtorch/BUILD"
)
diff --git a/docker/dist-build.sh b/docker/dist-build.sh
index 1655cf2c61..d3df8856c5 100755
--- a/docker/dist-build.sh
+++ b/docker/dist-build.sh
@@ -3,16 +3,20 @@
TOP_DIR=$(cd $(dirname $0); pwd)/..
if [[ -z "${USE_CXX11}" ]]; then
- BUILD_CMD="python setup.py bdist_wheel"
+ BUILD_CMD="python3 setup.py bdist_wheel"
else
- BUILD_CMD="python setup.py bdist_wheel --use-cxx11-abi"
+ BUILD_CMD="python3 setup.py bdist_wheel --use-cxx11-abi"
fi
cd ${TOP_DIR} \
&& mkdir -p dist && cd py \
- && pip install -r requirements.txt \
- && MAX_JOBS=1 LANG=en_US.UTF-8 LANGUAGE=en_US:en LC_ALL=en_US.UTF-8 \
- ${BUILD_CMD} $* || exit 1
+ && pip install -r requirements.txt
+
+# Symlink the path pyenv is using for python with the /opt directory for package sourcing
+ln -s "`pyenv which python | xargs dirname | xargs dirname`/lib/python$PYTHON_VERSION/site-packages" "/opt"
+
+# Build Torch-TRT
+MAX_JOBS=1 LANG=en_US.UTF-8 LANGUAGE=en_US:en LC_ALL=en_US.UTF-8 ${BUILD_CMD} $* || exit 1
pip3 install ipywidgets --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org
jupyter nbextension enable --py widgetsnbextension
From 054f5446f3d94fc217f44e8e1a73612af355d477 Mon Sep 17 00:00:00 2001
From: gs-olive <113141689+gs-olive@users.noreply.github.com>
Date: Wed, 12 Apr 2023 15:25:47 -0700
Subject: [PATCH 4/4] fix: Update python symlink path
- Add documentation regarding Python versioning
---
docker/Dockerfile | 6 +++---
docker/README.md | 3 ++-
docker/WORKSPACE.docker | 4 ++--
docker/WORKSPACE.ngc | 6 +++---
docker/dist-build.sh | 2 +-
5 files changed, 11 insertions(+), 10 deletions(-)
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 1e86bfc182..c2b262c13d 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -71,7 +71,7 @@ WORKDIR /workspace/torch_tensorrt/src
RUN cp ./docker/WORKSPACE.docker WORKSPACE
# Symlink the path pyenv is using for python with the /opt directory for package sourcing
-RUN ln -s "`pyenv which python | xargs dirname | xargs dirname`/lib/python$PYTHON_VERSION/site-packages" "/opt"
+RUN ln -s "`pyenv which python | xargs dirname | xargs dirname`/lib/python$PYTHON_VERSION/site-packages" "/opt/python3"
# This script builds both libtorchtrt bin/lib/include tarball and the Python wheel, in dist/
RUN bash ./docker/dist-build.sh
@@ -89,7 +89,7 @@ RUN pip install *.whl && rm -fr /workspace/torch_tensorrt/py/dist/* *.whl
WORKDIR /opt/torch_tensorrt
-ENV LD_LIBRARY_PATH /opt/site-packages/torch/lib:/opt/site-packages/torch_tensorrt/lib:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH}
-ENV PATH /opt/site-packages/torch_tensorrt/bin:${PATH}
+ENV LD_LIBRARY_PATH /opt/python3/site-packages/torch/lib:/opt/python3/site-packages/torch_tensorrt/lib:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH}
+ENV PATH /opt/python3/site-packages/torch_tensorrt/bin:${PATH}
CMD /bin/bash
diff --git a/docker/README.md b/docker/README.md
index a20a613553..2ed49c2b5f 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -4,7 +4,8 @@
* The `Dockerfile` currently uses Bazelisk to select the Bazel version, and uses the exact library versions of Torch and CUDA listed in dependencies.
* The desired versions of CUDNN and TensorRT must be specified as build-args, with major, minor, and patch versions as in: `--build-arg TENSORRT_VERSION=a.b.c --build-arg CUDNN_VERSION=x.y.z`
- * The desired base image be changed by explicitly setting a base image, as in `--build-arg BASE_IMG=nvidia/cuda:11.7.1-devel-ubuntu22.04`, though this is optional
+ * [**Optional**] The desired base image be changed by explicitly setting a base image, as in `--build-arg BASE_IMG=nvidia/cuda:11.7.1-devel-ubuntu22.04`, though this is optional
+ * [**Optional**] Additionally, the desired Python version can be changed by explicitly setting a version, as in `--build-arg PYTHON_VERSION=3.10`, though this is optional as well.
* This `Dockerfile` installs `pre-cxx11-abi` versions of Pytorch and builds Torch-TRT using `pre-cxx11-abi` libtorch as well.
diff --git a/docker/WORKSPACE.docker b/docker/WORKSPACE.docker
index 8a391c2075..dd5ff3cd9a 100755
--- a/docker/WORKSPACE.docker
+++ b/docker/WORKSPACE.docker
@@ -50,13 +50,13 @@ new_local_repository(
new_local_repository(
name = "libtorch",
- path = "/opt/site-packages/torch/",
+ path = "/opt/python3/site-packages/torch/",
build_file = "third_party/libtorch/BUILD"
)
new_local_repository(
name = "libtorch_pre_cxx11_abi",
- path = "/opt/site-packages/torch/",
+ path = "/opt/python3/site-packages/torch/",
build_file = "third_party/libtorch/BUILD"
)
diff --git a/docker/WORKSPACE.ngc b/docker/WORKSPACE.ngc
index 79bed14dc9..3e570c63f1 100755
--- a/docker/WORKSPACE.ngc
+++ b/docker/WORKSPACE.ngc
@@ -33,7 +33,7 @@ git_repository(
# This is currently used in pytorch NGC container CI testing.
local_repository(
name = "torch_tensorrt",
- path = "/usr/local/lib/python/site-packages/torch_tensorrt/"
+ path = "/opt/python3/site-packages/torch_tensorrt/"
)
# CUDA should be installed on the system locally
@@ -55,13 +55,13 @@ new_local_repository(
new_local_repository(
name = "libtorch",
- path = "/opt/site-packages/torch/",
+ path = "/opt/python3/site-packages/torch/",
build_file = "third_party/libtorch/BUILD"
)
new_local_repository(
name = "libtorch_pre_cxx11_abi",
- path = "/opt/site-packages/torch/",
+ path = "/opt/python3/site-packages/torch/",
build_file = "third_party/libtorch/BUILD"
)
diff --git a/docker/dist-build.sh b/docker/dist-build.sh
index d3df8856c5..ed5b271825 100755
--- a/docker/dist-build.sh
+++ b/docker/dist-build.sh
@@ -13,7 +13,7 @@ cd ${TOP_DIR} \
&& pip install -r requirements.txt
# Symlink the path pyenv is using for python with the /opt directory for package sourcing
-ln -s "`pyenv which python | xargs dirname | xargs dirname`/lib/python$PYTHON_VERSION/site-packages" "/opt"
+ln -s "`pyenv which python | xargs dirname | xargs dirname`/lib/python$PYTHON_VERSION/site-packages" "/opt/python3"
# Build Torch-TRT
MAX_JOBS=1 LANG=en_US.UTF-8 LANGUAGE=en_US:en LC_ALL=en_US.UTF-8 ${BUILD_CMD} $* || exit 1