feat(//docker): Adding CUDA11 based container for Ampere support

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

docker/Dockerfile.20.07

@@ -0,0 +1,31 @@
+FROM nvcr.io/nvidia/pytorch:20.07-py3
+
+RUN apt-get update && apt-get install curl gnupg  && rm -rf /var/lib/apt/lists/*
+
+RUN curl https://bazel.build/bazel-release.pub.gpg | apt-key add - && \
+    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 bazel-3.4.1 && rm -rf /var/lib/apt/lists/*
+RUN ln -s /usr/bin/bazel-3.4.1 /usr/bin/bazel
+
+RUN pip install notebook
+
+COPY . /workspace/TRTorch
+RUN rm /workspace/TRTorch/WORKSPACE
+COPY ./docker/WORKSPACE.cu11.docker  /workspace/TRTorch/WORKSPACE
+
+WORKDIR /workspace/TRTorch
+RUN bazel build //:libtrtorch --compilation_mode opt
+
+WORKDIR /workspace/TRTorch/py
+
+# Locale is not set by default
+RUN apt-get update && apt-get install -y locales ninja-build && rm -rf /var/lib/apt/lists/* && locale-gen en_US.UTF-8
+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 conda init bash
+
+WORKDIR /workspace/TRTorch/notebooks

docker/Dockerfile.docs

@@ -3,8 +3,8 @@ FROM nvcr.io/nvidia/tensorrt:20.03-py3
 RUN curl https://bazel.build/bazel-release.pub.gpg | apt-key add -
 RUN echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list
 
-RUN apt update && apt install bazel-3.3.1
-RUN ln -s /usr/bin/bazel-3.3.1 /usr/bin/bazel
+RUN apt update && apt install bazel-3.4.1
+RUN ln -s /usr/bin/bazel-3.4.1 /usr/bin/bazel
 
 
 COPY ./py/requirements.txt requirements.txt

docker/WORKSPACE.cu11.docker

@@ -0,0 +1,94 @@
+workspace(name = "TRTorch")
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
+
+git_repository(
+    name = "rules_python",
+    remote = "https://github.com/bazelbuild/rules_python.git",
+    commit = "4fcc24fd8a850bdab2ef2e078b1de337eea751a6",
+    shallow_since = "1589292086 -0400"
+)
+
+load("@rules_python//python:repositories.bzl", "py_repositories")
+py_repositories()
+
+load("@rules_python//python:pip.bzl", "pip_repositories", "pip3_import")
+pip_repositories()
+
+http_archive(
+    name = "rules_pkg",
+    url = "https://github.com/bazelbuild/rules_pkg/releases/download/0.2.4/rules_pkg-0.2.4.tar.gz",
+    sha256 = "4ba8f4ab0ff85f2484287ab06c0d871dcb31cc54d439457d28fd4ae14b18450a",
+)
+
+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"
+)
+
+# CUDA should be installed on the system locally
+new_local_repository(
+    name = "cuda",
+    path = "/usr/local/cuda-11.0/",
+    build_file = "@//third_party/cuda:BUILD",
+)
+
+new_local_repository(
+    name = "cublas",
+    path = "/usr",
+    build_file = "@//third_party/cublas:BUILD",
+)
+
+####################################################################################
+# Locally installed dependencies (use in cases of custom dependencies or aarch64)
+####################################################################################
+
+new_local_repository(
+    name = "libtorch",
+    path = "/opt/conda/lib/python3.6/site-packages/torch",
+    build_file = "third_party/libtorch/BUILD"
+)
+
+new_local_repository(
+    name = "libtorch_pre_cxx11_abi",
+    path = "/opt/conda/lib/python3.6/site-packages/torch",
+    build_file = "third_party/libtorch/BUILD"
+)
+
+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"
+)
+
+#########################################################################
+# Testing Dependencies (optional - comment out on aarch64)
+#########################################################################
+pip3_import(
+    name = "trtorch_py_deps",
+    requirements = "//py:requirements.txt"
+)
+
+load("@trtorch_py_deps//:requirements.bzl", "pip_install")
+pip_install()
+
+pip3_import(
+    name = "py_test_deps",
+    requirements = "//tests/py:requirements.txt"
+)
+
+load("@py_test_deps//:requirements.bzl", "pip_install")
+pip_install()
+

docker/WORKSPACE.docker

@@ -53,16 +53,16 @@ http_archive(
     name = "libtorch",
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/cu102/libtorch-cxx11-abi-shared-with-deps-1.5.1.zip"],
-    sha256 = "cf0691493d05062fe3239cf76773bae4c5124f4b039050dbdd291c652af3ab2a"
+    urls = ["https://download.pytorch.org/libtorch/cu102/libtorch-cxx11-abi-shared-with-deps-1.6.0.zip"],
+    sha256 = "fded948bd2dbee625cee33ebbd4843a69496729389e0200a90fbb667cdaeeb69"
 )
 
 http_archive(
     name = "libtorch_pre_cxx11_abi",
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
-    sha256 = "818977576572eadaf62c80434a25afe44dbaa32ebda3a0919e389dcbe74f8656",
-    urls = ["https://download.pytorch.org/libtorch/cu102/libtorch-shared-with-deps-1.5.1.zip"],
+    sha256 = "141bb229f4bbf905541096cf8705785e7b0c79e37ca1e5db9d372730b1b9abd7",
+    urls = ["https://download.pytorch.org/libtorch/cu102/libtorch-shared-with-deps-1.6.0.zip"],
 )
 
 ####################################################################################