feat(//docker): New base container to let master build in container

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

.dockerignore

@@ -60,4 +60,5 @@ examples/int8/qat/qat
 examples/int8/training/vgg16/data/*
 examples/int8/datasets/data/*
 env/**/*
-*.ts
+*.ts
+.nox/**/*

.github/workflows/docgen.yml

@@ -30,7 +30,7 @@ jobs:
         run: echo "::set-output name=sha::$(git rev-parse --short HEAD)"
       - name: Build Python Package
         run: |
-          cp docker/WORKSPACE.docs WORKSPACE
+          cp docker/WORKSPACE.docker WORKSPACE
           cd py
           python3 setup.py install
       - name: Generate New Docs

.github/workflows/linter.yml

@@ -34,7 +34,7 @@ jobs:
       - name: Run image
         run: |
           docker run -it -d --name cpplinter -e GITHUB_TOKEN=$GITHUB_TOKEN -v $GITHUB_WORKSPACE:/workspace -v $GITHUB_EVENT_PATH:/GITHUB_EVENT.json -w /workspace ghcr.io/nvidia/torch-tensorrt/docgen:latest
-          docker exec cpplinter bash -c "cp /workspace/docker/WORKSPACE.docs /workspace/WORKSPACE"
+          docker exec cpplinter bash -c "cp /workspace/docker/WORKSPACE.docker /workspace/WORKSPACE"
         env:
           GITHUB_TOKEN: ${{secrets.GITHUB_TOKEN}}
       - name: Install dependencies
@@ -71,7 +71,7 @@ jobs:
       - name: Run image
         run: |
           docker run -it -d --name pylinter -e GITHUB_TOKEN=$GITHUB_TOKEN -v $GITHUB_WORKSPACE:/workspace -v $GITHUB_EVENT_PATH:/GITHUB_EVENT.json -w /workspace ghcr.io/nvidia/torch-tensorrt/docgen:latest
-          docker exec pylinter bash -c "cp /workspace/docker/WORKSPACE.docs /workspace/WORKSPACE"
+          docker exec pylinter bash -c "cp /workspace/docker/WORKSPACE.docker /workspace/WORKSPACE"
         env:
           GITHUB_TOKEN: ${{secrets.GITHUB_TOKEN}}
       - name: Install dependencies

docker/Dockerfile

@@ -1,5 +1,5 @@
-ARG BASE=21.10
-ARG BASE_IMG=nvcr.io/nvidia/pytorch:${BASE}-py3
+ARG BASE=22.04
+ARG BASE_IMG=nvcr.io/nvidia/tensorrt:${BASE}-py3
 FROM ${BASE_IMG} as base
 
 FROM base as torch-tensorrt-builder-base
@@ -33,7 +33,7 @@ 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 Pythin wheel, in dist/
+# This script builds both libtorchtrt bin/lib/include tarball and the Python wheel, in dist/
 RUN ./docker/dist-build.sh
 
 FROM base as torch-tensorrt
@@ -44,20 +44,20 @@ RUN rm -rf /opt/pytorch/torch_tensorrt
 # copy source repo
 COPY . /workspace/torch_tensorrt
 COPY --from=torch-tensorrt-builder  /workspace/torch_tensorrt/src/py/dist/ .
-RUN conda init bash
+#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 pip3 install *.whl && rm -fr /workspace/torch_tensorrt/py/dist/* *.whl
 
-ENV LD_LIBRARY_PATH /opt/conda/lib/python3.8/site-packages/torch/lib:/opt/conda/lib/python3.8/site-packages/torch_tensorrt/lib:${LD_LIBRARY_PATH}
-ENV PATH /opt/conda/lib/python3.8/site-packages/torch_tensorrt/bin:${PATH}
+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}
+ENV PATH /usr/local/lib/python3.8/dist-packages/torch_tensorrt/bin:${PATH}
 #
 WORKDIR /workspace
-RUN mv /workspace/torch_tensorrt /opt/pytorch/torch_tensorrt
-RUN cp /opt/pytorch/torch_tensorrt/docker/WORKSPACE.docker /opt/pytorch/torch_tensorrt/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/pytorch/torch_tensorrt/notebooks /workspace/torch_tensorrt/notebooks
+RUN ln -s /opt/torch_tensorrt/notebooks /workspace/torch_tensorrt/notebooks
 
 CMD /bin/bash

docker/Dockerfile.ngc

@@ -0,0 +1,63 @@
+ARG BASE=21.10
+ARG BASE_IMG=nvcr.io/nvidia/pytorch:${BASE}-py3
+FROM ${BASE_IMG} as base
+
+FROM base as torch-tensorrt-builder-base
+
+# Removing any bazel or torch-tensorrt pre-installed from the base image
+RUN rm -rf /opt/pytorch/torch_tensorrt /usr/bin/bazel
+
+ARG ARCH="x86_64"
+ARG TARGETARCH="amd64"
+ARG BAZEL_VERSION=5.1.1
+
+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 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 Pythin wheel, in dist/
+RUN USE_CXX11=1 ./docker/dist-build.sh
+
+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 --from=torch-tensorrt-builder  /workspace/torch_tensorrt/src/py/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 pip3 install *.whl && rm -fr /workspace/torch_tensorrt/py/dist/* *.whl
+
+ENV LD_LIBRARY_PATH /opt/conda/lib/python3.8/site-packages/torch/lib:/opt/conda/lib/python3.8/site-packages/torch_tensorrt/lib:${LD_LIBRARY_PATH}
+ENV PATH /opt/conda/lib/python3.8/site-packages/torch_tensorrt/bin:${PATH}
+#
+WORKDIR /workspace
+RUN mv /workspace/torch_tensorrt /opt/pytorch/torch_tensorrt
+RUN cp /opt/pytorch/torch_tensorrt/docker/WORKSPACE.docker /opt/pytorch/torch_tensorrt/WORKSPACE
+RUN mkdir torch_tensorrt
+RUN ln -s /opt/pytorch/torch_tensorrt/notebooks /workspace/torch_tensorrt/notebooks
+
+CMD /bin/bash

docker/README.md

@@ -10,7 +10,7 @@ https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.
 
 ```
 # Build:
-DOCKER_BUILDKIT=1 docker build --build-arg BASE={PyTorch Base Container Version} -f docker/Dockerfile -t torch_tensorrt1.0:latest .
+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 \

docker/WORKSPACE.docker

@@ -4,93 +4,86 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
 
 http_archive(
-        name = "rules_python",
-        url = "https://github.com/bazelbuild/rules_python/releases/download/0.2.0/rules_python-0.2.0.tar.gz",
-        sha256 = "778197e26c5fbeb07ac2a2c5ae405b30f6cb7ad1f5510ea6fdac03bded96cc6f",
-    )
+    name = "rules_python",
+    sha256 = "778197e26c5fbeb07ac2a2c5ae405b30f6cb7ad1f5510ea6fdac03bded96cc6f",
+    url = "https://github.com/bazelbuild/rules_python/releases/download/0.2.0/rules_python-0.2.0.tar.gz",
+)
 
 load("@rules_python//python:pip.bzl", "pip_install")
 
 http_archive(
     name = "rules_pkg",
+    sha256 = "038f1caa773a7e35b3663865ffb003169c6a71dc995e39bf4815792f385d837d",
     urls = [
-    	"https://mirror.bazel.build/github.com/bazelbuild/rules_pkg/releases/download/0.4.0/rules_pkg-0.4.0.tar.gz",
-	"https://github.com/bazelbuild/rules_pkg/releases/download/0.4.0/rules_pkg-0.4.0.tar.gz",
+        "https://mirror.bazel.build/github.com/bazelbuild/rules_pkg/releases/download/0.4.0/rules_pkg-0.4.0.tar.gz",
+        "https://github.com/bazelbuild/rules_pkg/releases/download/0.4.0/rules_pkg-0.4.0.tar.gz",
     ],
-    sha256 = "038f1caa773a7e35b3663865ffb003169c6a71dc995e39bf4815792f385d837d",
 )
+
 load("@rules_pkg//:deps.bzl", "rules_pkg_dependencies")
+
 rules_pkg_dependencies()
 
 git_repository(
     name = "googletest",
-    remote = "https://github.com/google/googletest",
     commit = "703bd9caab50b139428cea1aaff9974ebee5742e",
-    shallow_since = "1570114335 -0400"
-)
-
-# External dependency for torch_tensorrt if you already have precompiled binaries.
-# This is currently used in pytorch NGC container CI testing.
-local_repository(
-    name = "torch_tensorrt",
-    path = "/opt/conda/lib/python3.8/site-packages/torch_tensorrt"
+    remote = "https://github.com/google/googletest",
+    shallow_since = "1570114335 -0400",
 )
 
 # CUDA should be installed on the system locally
 new_local_repository(
     name = "cuda",
-    path = "/usr/local/cuda",
+    path = "/usr/local/cuda/",
     build_file = "@//third_party/cuda:BUILD",
 )
 
 new_local_repository(
     name = "cublas",
-    build_file = "@//third_party/cublas:BUILD",
     path = "/usr",
+    build_file = "@//third_party/cublas:BUILD",
 )
 
-####################################################################################
-# Locally installed dependencies (use in cases of custom dependencies or aarch64)
-####################################################################################
+#############################################################################################################
+# Tarballs and fetched dependencies (default - use in cases when building from precompiled bin and tarballs)
+#############################################################################################################
 
-new_local_repository(
+http_archive(
     name = "libtorch",
-    path = "/opt/conda/lib/python3.8/site-packages/torch",
-    build_file = "third_party/libtorch/BUILD"
+    build_file = "@//third_party/libtorch:BUILD",
+    sha256 = "8d9e829ce9478db4f35bdb7943308cf02e8a2f58cf9bb10f742462c1d57bf287",
+    strip_prefix = "libtorch",
+    urls = ["https://download.pytorch.org/libtorch/cu113/libtorch-cxx11-abi-shared-with-deps-1.11.0%2Bcu113.zip"],
 )
 
-new_local_repository(
+http_archive(
     name = "libtorch_pre_cxx11_abi",
-    path = "/opt/conda/lib/python3.8/site-packages/torch",
-    build_file = "third_party/libtorch/BUILD"
+    build_file = "@//third_party/libtorch:BUILD",
+    sha256 = "90159ecce3ff451f3ef3f657493b6c7c96759c3b74bbd70c1695f2ea2f81e1ad",
+    strip_prefix = "libtorch",
+    urls = ["https://download.pytorch.org/libtorch/cu113/libtorch-shared-with-deps-1.11.0%2Bcu113.zip"],
 )
 
+####################################################################################
+# Locally installed dependencies (use in cases of custom dependencies or aarch64)
+####################################################################################
+
 new_local_repository(
     name = "cudnn",
     path = "/usr/",
     build_file = "@//third_party/cudnn/local:BUILD"
 )
 
 new_local_repository(
-    name = "tensorrt",
-    path = "/usr/",
-    build_file = "@//third_party/tensorrt/local:BUILD"
+   name = "tensorrt",
+   path = "/usr/",
+   build_file = "@//third_party/tensorrt/local:BUILD"
 )
 
 #########################################################################
 # Testing Dependencies (optional - comment out on aarch64)
 #########################################################################
-pip_install(
-    name = "torch_tensorrt_py_deps",
-    requirements = "//py:requirements.txt",
-)
-
-pip_install(
-    name = "py_test_deps",
-    requirements = "//tests/py:requirements.txt",
-)
-
 pip_install(
     name = "pylinter_deps",
     requirements = "//tools/linter:requirements.txt",
-)
+)