project-lakitu: cos-extensions: provide additional volume from cos-extensions for cos-gpu-installer

cos_gpu_installer v2: toolchain preloading/ caching

Empirically it has been observed that a significant amount of GPU driver
installation time is taken up during downloading and installing the
toolchain. This toolchain is a function of the build number of the cos
build running on the host machine.

One solution to cutting down on the gpu driver installation at run time
is by caching the toolchain at preload time. This can be achieved by
using the cos-extensions interface as follows:

At preload time, run cos-extensions install gpu -- --use-build-tools-cache.
This populates the cache.
At preload time, run rm -rf /var/lib/nvidia
At runtime, run cos-extensions install gpu -- --use-build-tools-cache to
install drivers using the populated cache.

This flow is achieved by providing an additional volume
/var/lib/cos-extensions -> /build from cos-extensions when invoking
cos-gpu-installer. /build is where the installer stores the toolchain
and it checks for the presence of the toolchain before downloading.

BUG=b/175236375
TEST=Tryjob
RELEASE_NOTE=Added the --use-build-tools-cache flag to `cos-extensions install gpu`.

Change-Id: I958fe6d259362a2ec33f5ba8e47b72fe3efc10c6
Reviewed-on: https://cos-review.googlesource.com/c/cos/overlays/board-overlays/+/19514
Reviewed-by: Roy Yang <royyang@google.com>
Tested-by: Arnav Kansal <rnv@google.com>
diff --git a/project-lakitu/app-admin/extensions-manager/extensions-manager-0.0.1-r6.ebuild b/project-lakitu/app-admin/extensions-manager/extensions-manager-0.0.1-r7.ebuild
similarity index 100%
rename from project-lakitu/app-admin/extensions-manager/extensions-manager-0.0.1-r6.ebuild
rename to project-lakitu/app-admin/extensions-manager/extensions-manager-0.0.1-r7.ebuild
diff --git a/project-lakitu/app-admin/extensions-manager/files/cos-extensions.sh b/project-lakitu/app-admin/extensions-manager/files/cos-extensions.sh
index 60b5163..6c7a73e 100755
--- a/project-lakitu/app-admin/extensions-manager/files/cos-extensions.sh
+++ b/project-lakitu/app-admin/extensions-manager/files/cos-extensions.sh
@@ -14,6 +14,7 @@
 readonly COS_GPU_INSTALLER="${COS_GPU_INSTALLER:-\
 gcr.io/cos-cloud/cos-gpu-installer:v2.0.6}"
 readonly OS_RELEASE="/etc/os-release"
+readonly EXTENSIONS_CACHE="/var/lib/cos-extensions"
 
 usage() {
   cat <<EOF
@@ -33,6 +34,14 @@
     install <extension> [-- -<version>]   install a COS extension. If no version
                                           is given, then the default version
                                           will be installed.
+
+Additional Description:
+    ${PROG_NAME} install gpu --use-build-tools-cache
+    The gpu extension can be invoked with a --use-build-tools-cache
+    optional argument that can be used to cache the toolchain for the installer.
+    Caching the toolchain carries the overhead of ~1GB disk space on the
+    stateful partition. It may also help to save time on downloading the
+    toolchain during the installation, after the first run.
 }
 EOF
   exit "${1}"
@@ -116,14 +125,33 @@
 
 run_gpu_installer() {
   check_arch
-  /usr/bin/docker run --rm \
-    --name="cos-gpu-installer" \
-    --privileged \
-    --net=host \
-    --pid=host \
-    --volume /dev:/dev \
-    --volume /:/root \
-    "${COS_GPU_INSTALLER}" "$@"
+  local use_build_cache=false
+
+  local installer_args=()
+  for i in "$@"; do
+    if [[ $i == '--use-build-tools-cache' ]]; then
+      use_build_cache=true
+    else
+      installer_args+=($i)
+    fi
+  done
+
+  local docker_args=(
+    --rm
+    --name="cos-gpu-installer"
+    --privileged
+    --net=host
+    --pid=host
+    --volume /dev:/dev
+    --volume /:/root
+  )
+
+  if [[ "$use_build_cache" = true ]]; then
+    docker_args+=(--volume ${EXTENSIONS_CACHE}/:/build/)
+  fi
+
+  /usr/bin/docker run "${docker_args[@]}" "${COS_GPU_INSTALLER}"\
+    "${installer_args[@]}"
 }
 
 main() {