Skip to content

Commit 970d775

Browse files
narendasannarendasan
committed
feat(//docker): Adding CUDA11 based container for Ampere support
Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
1 parent cc32c22 commit 970d775

File tree

4 files changed

+131
-6
lines changed

4 files changed

+131
-6
lines changed

Diff for: ‎docker/Dockerfile.20.07

+31
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
FROM nvcr.io/nvidia/pytorch:20.07-py3
2+
3+
RUN apt-get update && apt-get install curl gnupg && rm -rf /var/lib/apt/lists/*
4+
5+
RUN curl https://bazel.build/bazel-release.pub.gpg | apt-key add - && \
6+
echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list
7+
8+
RUN apt-get update && apt-get install bazel-3.4.1 && rm -rf /var/lib/apt/lists/*
9+
RUN ln -s /usr/bin/bazel-3.4.1 /usr/bin/bazel
10+
11+
RUN pip install notebook
12+
13+
COPY . /workspace/TRTorch
14+
RUN rm /workspace/TRTorch/WORKSPACE
15+
COPY ./docker/WORKSPACE.cu11.docker /workspace/TRTorch/WORKSPACE
16+
17+
WORKDIR /workspace/TRTorch
18+
RUN bazel build //:libtrtorch --compilation_mode opt
19+
20+
WORKDIR /workspace/TRTorch/py
21+
22+
# Locale is not set by default
23+
RUN apt-get update && apt-get install -y locales ninja-build && rm -rf /var/lib/apt/lists/* && locale-gen en_US.UTF-8
24+
ENV LANG en_US.UTF-8
25+
ENV LANGUAGE en_US:en
26+
ENV LC_ALL en_US.UTF-8
27+
RUN python3 setup.py install --use-cxx11-abi
28+
29+
RUN conda init bash
30+
31+
WORKDIR /workspace/TRTorch/notebooks

Diff for: ‎docker/Dockerfile.docs

+2-2
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ FROM nvcr.io/nvidia/tensorrt:20.03-py3
33
RUN curl https://bazel.build/bazel-release.pub.gpg | apt-key add -
44
RUN echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list
55

6-
RUN apt update && apt install bazel-3.3.1
7-
RUN ln -s /usr/bin/bazel-3.3.1 /usr/bin/bazel
6+
RUN apt update && apt install bazel-3.4.1
7+
RUN ln -s /usr/bin/bazel-3.4.1 /usr/bin/bazel
88

99

1010
COPY ./py/requirements.txt requirements.txt

Diff for: ‎docker/WORKSPACE.cu11.docker

+94
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
workspace(name = "TRTorch")
2+
3+
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
4+
load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
5+
6+
git_repository(
7+
name = "rules_python",
8+
remote = "https://github.com/bazelbuild/rules_python.git",
9+
commit = "4fcc24fd8a850bdab2ef2e078b1de337eea751a6",
10+
shallow_since = "1589292086 -0400"
11+
)
12+
13+
load("@rules_python//python:repositories.bzl", "py_repositories")
14+
py_repositories()
15+
16+
load("@rules_python//python:pip.bzl", "pip_repositories", "pip3_import")
17+
pip_repositories()
18+
19+
http_archive(
20+
name = "rules_pkg",
21+
url = "https://github.com/bazelbuild/rules_pkg/releases/download/0.2.4/rules_pkg-0.2.4.tar.gz",
22+
sha256 = "4ba8f4ab0ff85f2484287ab06c0d871dcb31cc54d439457d28fd4ae14b18450a",
23+
)
24+
25+
load("@rules_pkg//:deps.bzl", "rules_pkg_dependencies")
26+
rules_pkg_dependencies()
27+
28+
git_repository(
29+
name = "googletest",
30+
remote = "https://github.com/google/googletest",
31+
commit = "703bd9caab50b139428cea1aaff9974ebee5742e",
32+
shallow_since = "1570114335 -0400"
33+
)
34+
35+
# CUDA should be installed on the system locally
36+
new_local_repository(
37+
name = "cuda",
38+
path = "/usr/local/cuda-11.0/",
39+
build_file = "@//third_party/cuda:BUILD",
40+
)
41+
42+
new_local_repository(
43+
name = "cublas",
44+
path = "/usr",
45+
build_file = "@//third_party/cublas:BUILD",
46+
)
47+
48+
####################################################################################
49+
# Locally installed dependencies (use in cases of custom dependencies or aarch64)
50+
####################################################################################
51+
52+
new_local_repository(
53+
name = "libtorch",
54+
path = "/opt/conda/lib/python3.6/site-packages/torch",
55+
build_file = "third_party/libtorch/BUILD"
56+
)
57+
58+
new_local_repository(
59+
name = "libtorch_pre_cxx11_abi",
60+
path = "/opt/conda/lib/python3.6/site-packages/torch",
61+
build_file = "third_party/libtorch/BUILD"
62+
)
63+
64+
new_local_repository(
65+
name = "cudnn",
66+
path = "/usr/",
67+
build_file = "@//third_party/cudnn/local:BUILD"
68+
)
69+
70+
new_local_repository(
71+
name = "tensorrt",
72+
path = "/usr/",
73+
build_file = "@//third_party/tensorrt/local:BUILD"
74+
)
75+
76+
#########################################################################
77+
# Testing Dependencies (optional - comment out on aarch64)
78+
#########################################################################
79+
pip3_import(
80+
name = "trtorch_py_deps",
81+
requirements = "//py:requirements.txt"
82+
)
83+
84+
load("@trtorch_py_deps//:requirements.bzl", "pip_install")
85+
pip_install()
86+
87+
pip3_import(
88+
name = "py_test_deps",
89+
requirements = "//tests/py:requirements.txt"
90+
)
91+
92+
load("@py_test_deps//:requirements.bzl", "pip_install")
93+
pip_install()
94+

Diff for: ‎docker/WORKSPACE.docker

+4-4
Original file line numberDiff line numberDiff line change
@@ -53,16 +53,16 @@ http_archive(
5353
name = "libtorch",
5454
build_file = "@//third_party/libtorch:BUILD",
5555
strip_prefix = "libtorch",
56-
urls = ["https://download.pytorch.org/libtorch/cu102/libtorch-cxx11-abi-shared-with-deps-1.5.1.zip"],
57-
sha256 = "cf0691493d05062fe3239cf76773bae4c5124f4b039050dbdd291c652af3ab2a"
56+
urls = ["https://download.pytorch.org/libtorch/cu102/libtorch-cxx11-abi-shared-with-deps-1.6.0.zip"],
57+
sha256 = "fded948bd2dbee625cee33ebbd4843a69496729389e0200a90fbb667cdaeeb69"
5858
)
5959

6060
http_archive(
6161
name = "libtorch_pre_cxx11_abi",
6262
build_file = "@//third_party/libtorch:BUILD",
6363
strip_prefix = "libtorch",
64-
sha256 = "818977576572eadaf62c80434a25afe44dbaa32ebda3a0919e389dcbe74f8656",
65-
urls = ["https://download.pytorch.org/libtorch/cu102/libtorch-shared-with-deps-1.5.1.zip"],
64+
sha256 = "141bb229f4bbf905541096cf8705785e7b0c79e37ca1e5db9d372730b1b9abd7",
65+
urls = ["https://download.pytorch.org/libtorch/cu102/libtorch-shared-with-deps-1.6.0.zip"],
6666
)
6767

6868
####################################################################################

0 commit comments

Comments
 (0)