@@ -8,15 +8,20 @@ FROM base as torch-tensorrt-builder-base
8
8
9
9
RUN rm -rf /opt/torch-tensorrt /usr/bin/bazel
10
10
11
- RUN apt-get update && apt-get install --no-install-recommends -y curl gnupg
12
- RUN curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor > /etc/apt/trusted.gpg.d/bazel.gpg
13
- RUN echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list
11
+ ARG ARCH= "x86_64"
12
+ ARG TARGETARCH= "amd64"
13
+ ARG BAZEL_VERSION=4.0.0
14
14
15
- RUN apt-get update && apt-get install -y --no-install-recommends bazel-4.0.0
16
- RUN ln -s /usr/bin/bazel-4.0.0 /usr/bin/bazel
15
+ RUN [[ "$TARGETARCH" == "amd64" ]] && ARCH="x86_64" || ARCH="${TARGETARCH}" \
16
+ && wget -q https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION}/bazel-${BAZEL_VERSION}-linux-${ARCH} -O /usr/bin/bazel \
17
+ && chmod a+x /usr/bin/bazel
17
18
18
19
# Workaround for bazel expecting both static and shared versions, we only use shared libraries inside container
19
- RUN cp /usr/lib/x86_64-linux-gnu/libnvinfer.so /usr/lib/x86_64-linux-gnu/libnvinfer_static.a
20
+ RUN touch /usr/lib/$HOSTTYPE-linux-gnu/libnvinfer_static.a
21
+
22
+ RUN rm -rf /usr/local/cuda/lib* /usr/local/cuda/include \
23
+ && ln -sf /usr/local/cuda/targets/$HOSTTYPE-linux/lib /usr/local/cuda/lib64 \
24
+ && ln -sf /usr/local/cuda/targets/$HOSTTYPE-linux/include /usr/local/cuda/include
20
25
21
26
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
22
27
@@ -47,6 +52,7 @@ ENV PATH /opt/conda/lib/python3.8/site-packages/torch_tensorrt/bin:${PATH}
47
52
48
53
WORKDIR /workspace
49
54
RUN mv /workspace/torch_tensorrt /opt/pytorch/torch_tensorrt
55
+ # Make the notebooks visible to Jupyter in the current WORKDIR
50
56
RUN mkdir torch_tensorrt
51
57
RUN ln -s /opt/pytorch/torch_tensorrt/notebooks /workspace/torch_tensorrt/notebooks
52
58
0 commit comments