(//docker): Updated Dockerfile with NGC containers

Signed-off-by: Anurag Dixit <[email protected]>

Author: Anurag Dixit

docker/Dockerfile

@@ -0,0 +1,49 @@
+ARG BASE=21.08
+ARG BASE_IMG=nvcr.io/nvidia/pytorch:${BASE}-py3
+FROM ${BASE_IMG} as base
+
+FROM base as trtorch-builder-base
+
+# Removing any bazel or trtorch pre-installed from the base image
+
+RUN rm -rf /opt/pytorch/trtorch /usr/bin/bazel
+
+RUN apt-get update && apt-get install --no-install-recommends -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 --no-install-recommends 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 --no-install-recommends 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 pip3 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/dist-build.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+mkdir -p dist
+
+bazel build //:libtrtorch --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 uninstall -y trtorch && pip3 install dist/*.whl
+

docker/dist-test.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+pip3 install timm --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org
+# 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