Refactored Dockerfiles

README.md

@@ -206,6 +206,10 @@ docker run -it -v$(pwd)/..:/workspace/TRTorch build_trtorch_wheel /bin/bash /wor
 ```
 Python compilation expects using the tarball based compilation strategy from above.
 
+## Building Docker Image
+`Dockerfile` in `//docker` directory builds an image, based on latest NGC PyTorch container, with TRTorch binaries/libraries/headers installed in /opt/trtorch, and Python wheel installed locally.
+To specify different NGC PyTorch container release as the base, use `--build-arg BASE=21.06` for 21.06 etc.
+
 ## How do I add support for a new op...
 
 ### In TRTorch?

docker/Dockerfile

@@ -0,0 +1,44 @@
+ARG BASE=21.07
+FROM nvcr.io/nvidia/pytorch:${BASE}-py3 as base
+
+FROM base as trtorch-builder-base
+
+RUN apt-get update && apt-get install -y curl gnupg
+RUN curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor > /etc/apt/trusted.gpg.d/bazel.gpg
+RUN echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list
+
+RUN apt-get update && apt-get install -y bazel-4.0.0
+RUN ln -s /usr/bin/bazel-4.0.0 /usr/bin/bazel
+
+# Workaround for bazel expecting both static and shared versions, we only use shared libraries inside container
+RUN cp /usr/lib/x86_64-linux-gnu/libnvinfer.so /usr/lib/x86_64-linux-gnu/libnvinfer_static.a
+
+RUN apt-get update && apt-get install -y locales ninja-build && rm -rf /var/lib/apt/lists/* && locale-gen en_US.UTF-8
+
+FROM trtorch-builder-base as trtorch-builder
+
+COPY . /workspace/trtorch/src
+WORKDIR /workspace/trtorch/src
+RUN cp ./docker/WORKSPACE.cu.docker  WORKSPACE
+
+# This script builds both libtrtorch bin/lib/include tarball and the Pythin wheel, in dist/  
+RUN ./docker/dist-build.sh
+
+FROM base as trtorch
+
+# copy source repo
+COPY . /workspace/trtorch
+COPY --from=trtorch-builder  /workspace/trtorch/src/dist/ .
+
+RUN conda init bash
+
+RUN pip install ipywidgets --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org
+RUN jupyter nbextension enable --py widgetsnbextension
+
+RUN mkdir -p /opt/trtorch && tar xvf libtrtorch.tar.gz --strip-components 2 -C /opt/trtorch --exclude=LICENSE && pip3 install *.whl && rm -fr /workspace/trtorch/dist/*
+
+ENV LD_LIBRARY_PATH /opt/conda/lib/python3.8/site-packages/torch/lib:/opt/trtorch/lib:${LD_LIBRARY_PATH}
+ENV PATH /opt/trtorch/bin:${PATH}
+
+WORKDIR /workspace/trtorch/
+CMD /bin/bash

docker/Dockerfile.21.06

@@ -15,14 +15,13 @@ RUN pip install notebook
 
 FROM builder as trtorch
 
-COPY . /opt/trtorch
-RUN rm /opt/trtorch/WORKSPACE
-COPY ./docker/WORKSPACE.cu.docker  /opt/trtorch/WORKSPACE
+COPY . /opt/trtorch/src
+WORKDIR /opt/trtorch/src
+RUN cp ./docker/WORKSPACE.cu.docker  WORKSPACE
 
-WORKDIR /opt/trtorch
-RUN bazel build //:libtrtorch --compilation_mode opt
+RUN bazel build //:libtrtorch //:bin --compilation_mode opt
 
-WORKDIR /opt/trtorch/py
+WORKDIR /opt/trtorch/src/py
 
 RUN pip install ipywidgets --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org
 RUN jupyter nbextension enable --py widgetsnbextension
@@ -32,11 +31,15 @@ RUN apt-get update && apt-get install -y locales ninja-build && rm -rf /var/lib/
 ENV LANG en_US.UTF-8
 ENV LANGUAGE en_US:en
 ENV LC_ALL en_US.UTF-8
-RUN python3 setup.py install --use-cxx11-abi
+RUN MAX_JOBS=1 python3 setup.py install --use-cxx11-abi
 
 RUN conda init bash
 
-ENV LD_LIBRARY_PATH /opt/conda/lib/python3.8/site-packages/torch/lib:$LD_LIBRARY_PATH
+WORKDIR /opt/trtorch/src
+
+RUN cd ./bazel-bin && tar xvf libtrtorch.tar.gz -C /usr/local --strip-components 2 --exclude=LICENSE
+
+ENV LD_LIBRARY_PATH /opt/conda/lib/python3.8/site-packages/torch/lib:${LD_LIBRARY_PATH}
 
 WORKDIR /opt/trtorch/
 CMD /bin/bash

docker/dist-build.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+mkdir -p dist
+
+bazel build //:libtrtorch //:bin --compilation_mode opt
+
+cd py && MAX_JOBS=1 LANG=en_US.UTF-8 LANGUAGE=en_US:en LC_ALL=en_US.UTF-8 python3 setup.py bdist_wheel --use-cxx11-abi
+
+cd ..
+
+cp bazel-bin/libtrtorch.tar.gz dist/
+cp py/dist/* dist/
+
+pip3 install ipywidgets --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org
+jupyter nbextension enable --py widgetsnbextension
+
+pip3 install timm
+
+# test install
+mkdir -p /opt/trtorch && tar xvf dist/libtrtorch.tar.gz --strip-components 2 -C /opt/trtorch --exclude=LICENSE && pip3 install dist/*.whl
+

docker/dist-test.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+# Build and run unit tests
+cd tests/modules && python3 ./hub.py
+cd ../..
+
+bazel test //tests:tests //tests:python_api_tests --compilation_mode opt --jobs 4
+

py/requirements.txt

@@ -1,3 +1,3 @@
 -f https://download.pytorch.org/whl/torch_stable.html
-torch==1.9.0+cu111
+torch>=1.9.0+cu111
 pybind11==2.6.2

py/setup.py

@@ -203,7 +203,7 @@ def run(self):
             'trtorch/csrc/tensorrt_classes.cpp',
             'trtorch/csrc/register_tensorrt_classes.cpp',
         ],
-        library_dirs=[(dir_path + '/trtorch/lib/'), "/opt/conda/lib/python3.6/config-3.6m-x86_64-linux-gnu"],
+        library_dirs=[(dir_path + '/trtorch/lib/')],
         libraries=["trtorch"],
         include_dirs=[
             dir_path + "trtorch/csrc",
@@ -234,7 +234,7 @@ def run(self):
       long_description=long_description,
       ext_modules=ext_modules,
       install_requires=[
-          'torch>=1.9.0+cu111,<1.10.0',
+          'torch>=1.9.0+cu111,<1.11.0',
       ],
       setup_requires=[],
       cmdclass={

tests/modules/requirements.txt

@@ -1,3 +1,3 @@
 -f https://download.pytorch.org/whl/torch_stable.html
 timm==v0.4.12
-torch==1.9.0+cu111
+torch>=1.9.0+cu111