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/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/