Create a cuda 11.1 docker image for JAX.
PiperOrigin-RevId: 392667536
Change-Id: I782d77ec972b94a69d4a25b5fba9cf9f2eef87f8
diff --git a/tensorflow/tools/ci_build/Dockerfile.rbe.cuda11.1-cudnn8-ubuntu18.04-manylinux2010-multipython b/tensorflow/tools/ci_build/Dockerfile.rbe.cuda11.1-cudnn8-ubuntu18.04-manylinux2010-multipython
new file mode 100644
index 0000000..14eb34a
--- /dev/null
+++ b/tensorflow/tools/ci_build/Dockerfile.rbe.cuda11.1-cudnn8-ubuntu18.04-manylinux2010-multipython
@@ -0,0 +1,99 @@
+# Dockerfile to build a manylinux 2010 compliant cross-compiler.
+#
+# Builds a devtoolset gcc/libstdc++ that targets manylinux 2010 compatible
+# glibc (2.12) and system libstdc++ (4.4).
+#
+# To push a new version, run:
+# $ docker build -f Dockerfile.rbe.cuda11.1-cudnn8-ubuntu18.04-manylinux2010-multipython \
+# --tag "gcr.io/tensorflow-testing/nosla-cuda11.1-cudnn8-ubuntu18.04-manylinux2010-multipython" .
+# $ docker push gcr.io/tensorflow-testing/nosla-cuda11.1-cudnn8-ubuntu18.04-manylinux2010-multipython
+
+FROM nvidia/cuda:11.1.1-cudnn8-devel-ubuntu18.04 as devtoolset
+
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && apt-get install -y \
+ cpio \
+ file \
+ flex \
+ g++ \
+ make \
+ patch \
+ rpm2cpio \
+ unar \
+ wget \
+ xz-utils \
+ && \
+ rm -rf /var/lib/apt/lists/*
+
+ADD devtoolset/fixlinks.sh fixlinks.sh
+ADD devtoolset/build_devtoolset.sh build_devtoolset.sh
+ADD devtoolset/rpm-patch.sh rpm-patch.sh
+
+# Set up a sysroot for glibc 2.12 / libstdc++ 4.4 / devtoolset-7 in /dt7.
+RUN /build_devtoolset.sh devtoolset-7 /dt7
+# Set up a sysroot for glibc 2.12 / libstdc++ 4.4 / devtoolset-8 in /dt8.
+RUN /build_devtoolset.sh devtoolset-8 /dt8
+
+# TODO(klimek): Split up into two different docker images.
+FROM nvidia/cuda:11.1.1-cudnn8-devel-ubuntu18.04
+COPY --from=devtoolset /dt7 /dt7
+COPY --from=devtoolset /dt8 /dt8
+
+# Install TensorRT.
+RUN echo \
+ deb https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1604/x86_64 / \
+ > /etc/apt/sources.list.d/nvidia-ml.list \
+ && \
+ apt-get update && apt-get install -y \
+ libnvinfer-dev=7.2.2-1+cuda11.1 \
+ libnvinfer7=7.2.2-1+cuda11.1 \
+ libnvinfer-plugin-dev=7.2.2-1+cuda11.1 \
+ libnvinfer-plugin7=7.2.2-1+cuda11.1 \
+ && \
+ rm -rf /var/lib/apt/lists/*
+
+# Copy and run the install scripts.
+ARG DEBIAN_FRONTEND=noninteractive
+
+COPY install/install_bootstrap_deb_packages.sh /install/
+RUN /install/install_bootstrap_deb_packages.sh
+
+COPY install/install_deb_packages.sh /install/
+RUN /install/install_deb_packages.sh
+
+# Install additional packages needed for this image:
+# - dependencies to build Python from source
+# - patchelf, as it is required by auditwheel
+RUN apt-get update && apt-get install -y \
+ libbz2-dev \
+ libffi-dev \
+ libgdbm-dev \
+ libncurses5-dev \
+ libnss3-dev \
+ libreadline-dev \
+ libsqlite3-dev \
+ patchelf \
+ && \
+ rm -rf /var/lib/apt/lists/*
+
+COPY install/install_bazel.sh /install/
+RUN /install/install_bazel.sh
+
+COPY install/build_and_install_python.sh /install/
+RUN /install/build_and_install_python.sh "3.7.7"
+RUN /install/build_and_install_python.sh "3.8.2"
+RUN /install/build_and_install_python.sh "3.9.4"
+
+COPY install/install_pip_packages_by_version.sh /install/
+RUN /install/install_pip_packages_by_version.sh "/usr/local/bin/pip3.7"
+RUN /install/install_pip_packages_by_version.sh "/usr/local/bin/pip3.8"
+RUN /install/install_pip_packages_by_version.sh "/usr/local/bin/pip3.9"
+
+ENV CLANG_VERSION="r7f6f9f4cf966c78a315d15d6e913c43cfa45c47c"
+COPY install/install_latest_clang.sh /install/
+RUN /install/install_latest_clang.sh
+
+# TensorRT 7 for CUDA 11.1 is compatible with CUDA 11.2, but requires
+# libnvrtc.so.11.1. See https://github.com/NVIDIA/TensorRT/issues/1064.
+# TODO(b/187962120): Remove when upgrading to TensorRT 8.
+ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/cuda-11.1/lib64"
diff --git a/tensorflow/workspace3.bzl b/tensorflow/workspace3.bzl
index 2b4c4dc..020b73c 100644
--- a/tensorflow/workspace3.bzl
+++ b/tensorflow/workspace3.bzl
@@ -17,11 +17,11 @@
http_archive(
name = "tf_toolchains",
- sha256 = "072753865000f81b5686edd273ae74d35361cfa0b16bb1b3400c25d0908f4863",
- strip_prefix = "toolchains-1.2.6",
+ sha256 = "8613ab9da262fd30d101550eab5780948db7471c7809cbac15a9e9a3143447f2",
+ strip_prefix = "toolchains-1.2.7",
urls = [
- "http://mirror.tensorflow.org/github.com/tensorflow/toolchains/archive/v1.2.6.tar.gz",
- "https://github.com/tensorflow/toolchains/archive/v1.2.6.tar.gz",
+ "http://mirror.tensorflow.org/github.com/tensorflow/toolchains/archive/v1.2.7.tar.gz",
+ "https://github.com/tensorflow/toolchains/archive/v1.2.7.tar.gz",
],
)
