From 7b1707bb7dd2589f551360f994d0866cdff22876 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] 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..2ed49ff1c2 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 /usr/local/lib/ 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..4ba69f1b31 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 /usr/local/lib/ 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