Update Vulkan SDK within CI docker images (#3825)

This commit changes to download the Vulkan SDK tarball instead of
installing it from APT repo so that we can still get it for the
RBE toolchain. LunarG APT repo dropped support for Ubuntu 16.04
for recent SDK versions.

This commit also creates a base Vulkan SDK image so that other
docker images can copy the SDK from it.

This commit also pin the NVIDIA driver and `tf-nightly` to a specific
version to increase stability. A previous pull request
https://github.com/google/iree/pull/3754 with mostly the same
change was reverted due to failures on TensorFlow integration
tests. That should be fixed now by pinning down the `tf-nightly`
version.
diff --git a/build_tools/docker/bazel-tensorflow-nvidia/Dockerfile b/build_tools/docker/bazel-tensorflow-nvidia/Dockerfile
index a9b4cf4..adb2e60 100644
--- a/build_tools/docker/bazel-tensorflow-nvidia/Dockerfile
+++ b/build_tools/docker/bazel-tensorflow-nvidia/Dockerfile
@@ -18,4 +18,4 @@
 FROM gcr.io/iree-oss/bazel-tensorflow-vulkan AS final
 
 RUN apt-get update \
-  && DEBIAN_FRONTEND=noninteractive apt-get install -y vulkan-sdk nvidia-driver-440
+  && DEBIAN_FRONTEND=noninteractive apt-get install -y libnvidia-gl-440=440.100-0ubuntu0.18.04.1
diff --git a/build_tools/docker/bazel-tensorflow-vulkan/Dockerfile b/build_tools/docker/bazel-tensorflow-vulkan/Dockerfile
index 849a32c..05b7209 100644
--- a/build_tools/docker/bazel-tensorflow-vulkan/Dockerfile
+++ b/build_tools/docker/bazel-tensorflow-vulkan/Dockerfile
@@ -13,23 +13,21 @@
 # limitations under the License.
 
 # A base image for building IREE with TensorFlow integrations using Bazel and
-# running Vulkan tests. Requires a child image to provide a Vulkan ICD.
+# running Vulkan tests. This image provides the Vulkan SDK. Requires a child
+# image to provide a Vulkan ICD.
 
 FROM gcr.io/iree-oss/bazel-tensorflow AS final
 
-RUN apt-get update && apt-get install -y wget
+ARG VULKAN_SDK_VERSION=1.2.154.0
 
-ARG VULKAN_SDK_VERSION=1.2.141
+COPY --from=gcr.io/iree-oss/vulkan /opt/vulkan-sdk/ /opt/vulkan-sdk/
 
-# Disable apt-key parse waring. If someone knows how to do whatever the "proper"
-# thing is then feel free. The warning complains about parsing apt-key output,
-# which we're not even doing.
-ARG APT_KEY_DONT_WARN_ON_DANGEROUS_USAGE=1
+ENV VULKAN_SDK="/opt/vulkan-sdk/${VULKAN_SDK_VERSION}/x86_64"
 
-RUN wget -qO - http://packages.lunarg.com/lunarg-signing-key-pub.asc \
-    | apt-key add - \
-  && wget -qO \
-    "/etc/apt/sources.list.d/lunarg-vulkan-${VULKAN_SDK_VERSION?}-bionic.list" \
-    "http://packages.lunarg.com/vulkan/${VULKAN_SDK_VERSION?}/lunarg-vulkan-${VULKAN_SDK_VERSION?}-bionic.list" \
-  && apt-get update \
-  && apt-get install -y vulkan-sdk
+ENV PATH="${VULKAN_SDK}/bin:$PATH"
+
+# Symlink the Vulkan loader to a system library directory. This is needed to
+# allow Vulkan applications to find the Vulkan loader. It also avoids using
+# LD_LIBRARY_PATH, which is not supported well by Docker.
+RUN ln -s "${VULKAN_SDK}/lib/libvulkan.so" /usr/lib/x86_64-linux-gnu/ \
+  && ln -s "${VULKAN_SDK}/lib/libvulkan.so.1" /usr/lib/x86_64-linux-gnu/
diff --git a/build_tools/docker/bazel-tensorflow/Dockerfile b/build_tools/docker/bazel-tensorflow/Dockerfile
index a9dd7fb..e1801c1 100644
--- a/build_tools/docker/bazel-tensorflow/Dockerfile
+++ b/build_tools/docker/bazel-tensorflow/Dockerfile
@@ -17,4 +17,4 @@
 FROM gcr.io/iree-oss/bazel-python AS final
 
 # Install tensorflow.
-RUN python3 -m pip install tf-nightly
+RUN python3 -m pip install tf-nightly==2.4.0.dev20200902
diff --git a/build_tools/docker/cmake-python-nvidia/Dockerfile b/build_tools/docker/cmake-python-nvidia/Dockerfile
index cce8d2f..e9559d4 100644
--- a/build_tools/docker/cmake-python-nvidia/Dockerfile
+++ b/build_tools/docker/cmake-python-nvidia/Dockerfile
@@ -28,4 +28,4 @@
 FROM gcr.io/iree-oss/cmake-python-vulkan AS final
 
 RUN apt-get update \
-  && apt-get install -y nvidia-driver-440
+  && DEBIAN_FRONTEND=noninteractive apt-get install -y libnvidia-gl-440=440.100-0ubuntu0.18.04.1
diff --git a/build_tools/docker/cmake-python-vulkan/Dockerfile b/build_tools/docker/cmake-python-vulkan/Dockerfile
index a20221b..af95afc 100644
--- a/build_tools/docker/cmake-python-vulkan/Dockerfile
+++ b/build_tools/docker/cmake-python-vulkan/Dockerfile
@@ -12,27 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# A base image for building IREE using CMake and running Vulkan tests. Requires
-# a child image to provide a Vulkan ICD.
+# A base image for building IREE using CMake and running Vulkan tests.
+# This image provides the Vulkan SDK. Requires a child image to provide
+# a Vulkan ICD.
 
 FROM gcr.io/iree-oss/cmake-python AS final
 
-# It would be nice to have a separate install vulkan image and copy from that,
-# but I don't know all the files that installing the vulkan-sdk adds.
-RUN apt-get update && apt-get install -y wget
+ARG VULKAN_SDK_VERSION=1.2.154.0
 
-ARG VULKAN_SDK_VERSION=1.2.141
-# Disable apt-key parse waring. If someone knows how to do whatever the "proper"
-# thing is then feel free. The warning complains about parsing apt-key output,
-# which we're not even doing.
-ARG APT_KEY_DONT_WARN_ON_DANGEROUS_USAGE=1
+COPY --from=gcr.io/iree-oss/vulkan /opt/vulkan-sdk/ /opt/vulkan-sdk/
 
-RUN wget -qO - http://packages.lunarg.com/lunarg-signing-key-pub.asc \
-    | apt-key add - \
-  && wget -qO \
-    "/etc/apt/sources.list.d/lunarg-vulkan-${VULKAN_SDK_VERSION?}-bionic.list" \
-    "http://packages.lunarg.com/vulkan/${VULKAN_SDK_VERSION?}/lunarg-vulkan-${VULKAN_SDK_VERSION?}-bionic.list" \
-  && apt-get update \
-  && apt-get install -y vulkan-sdk
+ENV VULKAN_SDK="/opt/vulkan-sdk/${VULKAN_SDK_VERSION}/x86_64"
 
-RUN rm /usr/bin/wget
+ENV PATH="${VULKAN_SDK}/bin:$PATH"
+
+# Symlink the Vulkan loader to a system library directory. This is needed to
+# allow Vulkan applications to find the Vulkan loader. It also avoids using
+# LD_LIBRARY_PATH, which is not supported well by Docker.
+RUN ln -s "${VULKAN_SDK}/lib/libvulkan.so" /usr/lib/x86_64-linux-gnu/ \
+  && ln -s "${VULKAN_SDK}/lib/libvulkan.so.1" /usr/lib/x86_64-linux-gnu/
diff --git a/build_tools/docker/manage_images.py b/build_tools/docker/manage_images.py
index e8c7c8b..54773a0 100755
--- a/build_tools/docker/manage_images.py
+++ b/build_tools/docker/manage_images.py
@@ -54,16 +54,17 @@
     'bazel-tensorflow': ['bazel-python'],
     'bazel-tensorflow-nvidia': ['bazel-tensorflow-vulkan'],
     'bazel-tensorflow-swiftshader': ['bazel-tensorflow-vulkan', 'swiftshader'],
-    'bazel-tensorflow-vulkan': ['bazel-tensorflow'],
+    'bazel-tensorflow-vulkan': ['bazel-tensorflow', 'vulkan'],
     'cmake': ['base', 'util'],
     'cmake-android': ['cmake', 'util'],
     'cmake-python': ['cmake'],
     'cmake-python-nvidia': ['cmake-python-vulkan'],
     'cmake-python-swiftshader': ['cmake-python-vulkan', 'swiftshader'],
-    'cmake-python-vulkan': ['cmake-python'],
-    'rbe-toolchain': [],
+    'cmake-python-vulkan': ['cmake-python', 'vulkan'],
+    'rbe-toolchain': ['vulkan'],
     'swiftshader': ['cmake'],
     'util': [],
+    'vulkan': ['util'],
 }
 
 IMAGES_TO_DEPENDENT_IMAGES = {k: [] for k in IMAGES_TO_DEPENDENCIES}
diff --git a/build_tools/docker/rbe-toolchain/Dockerfile b/build_tools/docker/rbe-toolchain/Dockerfile
index ffb0f0f..f3fef55 100644
--- a/build_tools/docker/rbe-toolchain/Dockerfile
+++ b/build_tools/docker/rbe-toolchain/Dockerfile
@@ -79,25 +79,20 @@
   && python3 -m pip install numpy
 
 ######################## Vulkan SDK ############################################
-ARG VULKAN_SDK_VERSION=1.2.141
 
-# Disable apt-key parse waring. If someone knows how to do whatever the "proper"
-# thing is then feel free. The warning complains about parsing apt-key output,
-# which we're not even doing.
-ARG APT_KEY_DONT_WARN_ON_DANGEROUS_USAGE=1
+ARG VULKAN_SDK_VERSION=1.2.154.0
 
-RUN apt-get update \
-  && apt-get install apt-transport-https
+COPY --from=gcr.io/iree-oss/vulkan /opt/vulkan-sdk/ /opt/vulkan-sdk/
 
-# Note that this image is based on Ubuntu 16.04 (xenial) as opposed to
-# Ubuntu 18.04 (bionic), which we use for our other images.
-RUN wget -qO - http://packages.lunarg.com/lunarg-signing-key-pub.asc \
-    | apt-key add - \
-  && wget -qO \
-    "/etc/apt/sources.list.d/lunarg-vulkan-${VULKAN_SDK_VERSION?}-xenial.list" \
-    "http://packages.lunarg.com/vulkan/${VULKAN_SDK_VERSION?}/lunarg-vulkan-${VULKAN_SDK_VERSION?}-xenial.list" \
-  && apt-get update \
-  && apt-get install -y vulkan-sdk
+ENV VULKAN_SDK="/opt/vulkan-sdk/${VULKAN_SDK_VERSION}/x86_64"
+
+ENV PATH="${VULKAN_SDK}/bin:$PATH"
+
+# Symlink the Vulkan loader to a system library directory. This is needed to
+# allow Vulkan applications to find the Vulkan loader. It also avoids using
+# LD_LIBRARY_PATH, which is not supported well by Docker.
+RUN ln -s "${VULKAN_SDK}/lib/libvulkan.so" /usr/lib/x86_64-linux-gnu/ \
+  && ln -s "${VULKAN_SDK}/lib/libvulkan.so.1" /usr/lib/x86_64-linux-gnu/
 
 ######################## Swiftshader ###########################################
 COPY --from=install-swiftshader /swiftshader /swiftshader
diff --git a/build_tools/docker/vulkan/Dockerfile b/build_tools/docker/vulkan/Dockerfile
new file mode 100644
index 0000000..24a8914
--- /dev/null
+++ b/build_tools/docker/vulkan/Dockerfile
@@ -0,0 +1,25 @@
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+FROM gcr.io/iree-oss/util AS base
+
+ARG VULKAN_SDK_VERSION=1.2.154.0
+
+RUN wget -q \
+  "https://sdk.lunarg.com/sdk/download/${VULKAN_SDK_VERSION?}/linux/vulkansdk-linux-${VULKAN_SDK_VERSION?}.tar.gz" \
+  -O "/tmp/vulkansdk-linux-x86_64-${VULKAN_SDK_VERSION}.tar.gz"
+
+RUN mkdir -p /opt/vulkan-sdk
+
+RUN tar -xf /tmp/vulkansdk-linux-x86_64-$VULKAN_SDK_VERSION.tar.gz -C /opt/vulkan-sdk