[NFC] Refactor docker scripting to generalize away from Kokoro

build_tools/kokoro/gcp_ubuntu/docker_common.sh → build_tools/docker/docker_run.sh

@@ -4,22 +4,19 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-# Functions for setting up Docker containers to run on Kokoro
+set -euo pipefail
 
-# Sets up files and environment to enable running all our Kokoro docker scripts.
+# Sets up files and environment to enable running scripts in docker.
 # In particular, does some shenanigans to enable running with the current user.
 # Some of this setup is only strictly necessary for Bazel, but it doesn't hurt
 # for anything else.
-# Requires that KOKORO_ROOT and KOKORO_ARTIFACTS_DIR have been set
-# Sets the environment variable DOCKER_RUN_ARGS to be used by subsequent
-# `docker run` invocations.
-function docker_setup() {
+# Requires that DOCKER_WORKDIR and DOCKER_TMPDIR have been set
+function docker_run() {
     # Make the source repository available and launch containers in that
     # directory.
-    local workdir="${KOKORO_ARTIFACTS_DIR?}/github/iree"
     DOCKER_RUN_ARGS=(
-      --volume="${workdir?}:${workdir?}"
-      --workdir="${workdir?}"
+      --volume="${DOCKER_WORKDIR}:${DOCKER_WORKDIR}"
+      --workdir="${DOCKER_WORKDIR}"
     )
 
     # Delete the container after the run is complete.
@@ -48,7 +45,7 @@ function docker_setup() {
     # such that they don't contain the information about normal users and we
     # want these scripts to be runnable locally for debugging.
     # Instead we dump the results of `getent` to some fake files.
-    local fake_etc_dir="${KOKORO_ROOT?}/fake_etc"
+    local fake_etc_dir="${DOCKER_TMPDIR}/fake_etc"
     mkdir -p "${fake_etc_dir?}"
 
     local fake_group="${fake_etc_dir?}/group"
@@ -71,11 +68,12 @@ function docker_setup() {
     # for two reasons:
     #   1. We probably don't want Docker to just write into the user's home
     #      directory when running locally.
-    #   2. When running with Kokoro, we mount a local scratch SSD to KOKORO_ROOT
-    #      whereas the home directory is on the persistent SSD boot disk. It
-    #      turns out that makes a huge difference in performance for Bazel
-    #      because it is IO bound at 64 cores..
-    local fake_home_dir="${KOKORO_ROOT?}/fake_home"
+    #   2. This allows us to control the device the home directory is in. Bazel
+    #      tends to be IO bound at even moderate levels of CPU parallelism and
+    #      the difference between a persistent SSD and a local scratch SSD can
+    #      be huge. In particular, Kokoro has the home directory on the former
+    #      and the work directory on the latter.
+    local fake_home_dir="${DOCKER_TMPDIR}/fake_home"
     mkdir -p "${fake_home_dir}"
 
     DOCKER_RUN_ARGS+=(
@@ -87,4 +85,8 @@ function docker_setup() {
     DOCKER_RUN_ARGS+=(
       --volume="${HOME?}/.config/gcloud:${HOME?}/.config/gcloud:ro"
     )
+
+    docker run "${DOCKER_RUN_ARGS[@]}" "$@"
 }
+
+docker_run "$@"

build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/core/build_kokoro.sh

@@ -18,16 +18,11 @@ set -o pipefail
 # Print the UTC time when set -x is on
 export PS4='[$(date -u "+%T %Z")] '
 
-source "${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_common.sh"
-
-# Sets DOCKER_RUN_ARGS
-docker_setup
-
 # This doesn't really need everything in the frontends image, but we want the
 # cache to be shared with the integrations build (no point building LLVM twice)
 # and the cache key is the docker container it's run in (to ensure correct cache
 # hits).
-docker run "${DOCKER_RUN_ARGS[@]?}" \
+"${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_run.sh" \
   gcr.io/iree-oss/frontends-swiftshader@sha256:41e516b8c1b432e3c02896c4bf4b7f06df6a67371aa167b88767b8d4d2018ea6 \
   build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/core/build.sh
 

build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-swiftshader/build_kokoro.sh

@@ -18,12 +18,7 @@ set -o pipefail
 # Print the UTC time when set -x is on
 export PS4='[$(date -u "+%T %Z")] '
 
-source "${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_common.sh"
-
-# Sets DOCKER_RUN_ARGS
-docker_setup
-
-docker run "${DOCKER_RUN_ARGS[@]?}" \
+"${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_run.sh" \
   gcr.io/iree-oss/frontends-swiftshader@sha256:41e516b8c1b432e3c02896c4bf4b7f06df6a67371aa167b88767b8d4d2018ea6 \
   build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-swiftshader/build.sh
 

build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-turing/build_kokoro.sh

@@ -18,17 +18,12 @@ set -o pipefail
 # Print the UTC time when set -x is on
 export PS4='[$(date -u "+%T %Z")] '
 
-source "${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_common.sh"
-
 # Print NVIDIA GPU information inside the VM
 dmesg | grep NVRM
 dpkg -l | grep nvidia
 nvidia-smi || true
 
-# Sets DOCKER_RUN_ARGS
-docker_setup
-
-docker run "${DOCKER_RUN_ARGS[@]?}" \
+"${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_run.sh" \
   --gpus all \
   gcr.io/iree-oss/frontends-nvidia@sha256:311ae679618bd248692522ea243ff73c8205725c95b6d915dea3a74dca9929ed \
   build_tools/kokoro/gcp_ubuntu/cmake-bazel/linux/x86-turing/build.sh

build_tools/kokoro/gcp_ubuntu/cmake/android/arm64-v8a/build_kokoro.sh

@@ -11,21 +11,14 @@
 # Requires the environment variables KOKORO_ROOT and KOKORO_ARTIFACTS_DIR, which
 # are set by Kokoro.
 
-set -x
-set -e
-set -o pipefail
+set -xeuo pipefail
 
 # Print the UTC time when set -x is on
 export PS4='[$(date -u "+%T %Z")] '
 
-source "${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_common.sh"
-
-# Sets DOCKER_RUN_ARGS
-docker_setup
-
 # Need to use frontends image (which also has Android toolchain) to build the
 # TFLite compiler for generating benchmarks.
-docker run "${DOCKER_RUN_ARGS[@]?}" \
+"${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_run.sh" \
   gcr.io/iree-oss/frontends@sha256:908d7e0ba3e3a6d448b3c07b39cc0f786a120dd2ced8653993f62fa933e882f7 \
   build_tools/kokoro/gcp_ubuntu/cmake/android/build.sh arm64-v8a
 

build_tools/kokoro/gcp_ubuntu/cmake/baremetal/riscv32/build_kokoro.sh

@@ -18,12 +18,7 @@ set -o pipefail
 # Print the UTC time when set -x is on
 export PS4='[$(date -u "+%T %Z")] '
 
-source "${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_common.sh"
-
-# Sets DOCKER_RUN_ARGS
-docker_setup
-
-docker run "${DOCKER_RUN_ARGS[@]?}" \
+"${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_run.sh" \
   gcr.io/iree-oss/riscv@sha256:e566f054ff1b1d8be61459ce4789dd944b70e85c8939a4d3b7331ab519d8db4c \
   build_tools/kokoro/gcp_ubuntu/cmake/baremetal/riscv32/build.sh
 

build_tools/kokoro/gcp_ubuntu/cmake/linux/riscv64/build_kokoro.sh

@@ -18,12 +18,7 @@ set -o pipefail
 # Print the UTC time when set -x is on
 export PS4='[$(date -u "+%T %Z")] '
 
-source "${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_common.sh"
-
-# Sets DOCKER_RUN_ARGS
-docker_setup
-
-docker run "${DOCKER_RUN_ARGS[@]?}" \
+"${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_run.sh" \
   gcr.io/iree-oss/riscv@sha256:e566f054ff1b1d8be61459ce4789dd944b70e85c8939a4d3b7331ab519d8db4c \
   build_tools/kokoro/gcp_ubuntu/cmake/linux/riscv64/build.sh
 

build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader-asan/build_kokoro.sh

@@ -18,12 +18,7 @@ set -o pipefail
 # Print the UTC time when set -x is on
 export PS4='[$(date -u "+%T %Z")] '
 
-source "${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_common.sh"
-
-# Sets DOCKER_RUN_ARGS
-docker_setup
-
-docker run "${DOCKER_RUN_ARGS[@]?}" \
+"${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_run.sh" \
   gcr.io/iree-oss/swiftshader@sha256:542de3bac7f6173add3651f75c4c681f4a8e2d23815332dd82afe85c7d598445 \
   build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader-asan/build.sh
 

build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader/build_kokoro.sh

@@ -18,12 +18,7 @@ set -o pipefail
 # Print the UTC time when set -x is on
 export PS4='[$(date -u "+%T %Z")] '
 
-source "${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_common.sh"
-
-# Sets DOCKER_RUN_ARGS
-docker_setup
-
-docker run "${DOCKER_RUN_ARGS[@]?}" \
+"${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_run.sh" \
   gcr.io/iree-oss/swiftshader@sha256:542de3bac7f6173add3651f75c4c681f4a8e2d23815332dd82afe85c7d598445 \
   build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader/build.sh
 

build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-turing/build_kokoro.sh

@@ -18,17 +18,12 @@ set -o pipefail
 # Print the UTC time when set -x is on
 export PS4='[$(date -u "+%T %Z")] '
 
-source "${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_common.sh"
-
 # Print NVIDIA GPU information inside the VM
 dmesg | grep NVRM
 dpkg -l | grep nvidia
 nvidia-smi || true
 
-# Sets DOCKER_RUN_ARGS
-docker_setup
-
-docker run "${DOCKER_RUN_ARGS[@]?}" \
+"${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/kokoro/gcp_ubuntu/docker_run.sh" \
   --gpus all \
   gcr.io/iree-oss/nvidia@sha256:8d7211e10199136482428b845b312169533f9119d2b9aeffb1ddd942b9e444a0 \
   build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-turing/build.sh

build_tools/kokoro/gcp_ubuntu/docker_run.sh

@@ -0,0 +1,16 @@
+# Copyright 2020 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# Runs docker configured for usage with Kokoro actions, translating Kokoro
+# environment variables into generic ones and then invoking the generic
+# docker_run script.
+
+set -euo pipefail
+
+export DOCKER_WORKDIR="${KOKORO_ARTIFACTS_DIR}/github/iree"
+export DOCKER_TMPDIR="${KOKORO_ROOT}"
+
+"${KOKORO_ARTIFACTS_DIR?}/github/iree/build_tools/docker/docker_run.sh" "$@"