Merge "Adds gpu driver setup to host image build scripts"

2 files changed, 51 insertions, 4 deletions

diff --git a/tools/create_base_image_gce.sh b/tools/create_base_image_gce.sh
index 677ae263..09443e37 100755
--- a/tools/create_base_image_gce.sh
+++ b/tools/create_base_image_gce.sh
@@ -53,14 +53,57 @@ sudo chroot /mnt/image /usr/bin/apt update
 sudo chroot /mnt/image /usr/bin/apt install -y "${tmp_debs[@]}"
 # install tools dependencies
 sudo chroot /mnt/image /usr/bin/apt install -y openjdk-11-jre
-sudo chroot /mnt/image /usr/bin/apt install -y unzip bzip2
+sudo chroot /mnt/image /usr/bin/apt install -y unzip bzip2 lzop
 sudo chroot /mnt/image /usr/bin/apt install -y aapt
 
 sudo chroot /mnt/image /usr/bin/find /home -ls
+
+
+# Install GPU driver dependencies
+sudo chroot /mnt/image /usr/bin/apt install -y gcc
+sudo chroot /mnt/image /usr/bin/apt install -y linux-source
+sudo chroot /mnt/image /usr/bin/apt install -y linux-headers-`uname -r`
+sudo chroot /mnt/image /usr/bin/apt install -y make
+
+# Download the latest GPU driver installer
+gsutil cp \
+  $(gsutil ls gs://nvidia-drivers-us-public/GRID/GRID*/*-Linux-x86_64-*.run \
+    | sort \
+    | tail -n 1) \
+  /mnt/image/tmp/nvidia-driver-installer.run
+
+# Make GPU driver installer executable
+chmod +x /mnt/image/tmp/nvidia-driver-installer.run
+
+# Install the latest GPU driver with default options and the dispatch libs
+sudo chroot /mnt/image /tmp/nvidia-driver-installer.run \
+  --silent \
+  --install-libglvnd
+
+# Cleanup after install
+rm /mnt/image/tmp/nvidia-driver-installer.run
+
+# Verify
+query_nvidia() {
+  sudo chroot /mnt/image nvidia-smi --format=csv,noheader --query-gpu="$@"
+}
+
+if [[ $(query_nvidia "count") != "1" ]]; then
+  echo "Failed to detect GPU."
+  exit 1
+fi
+
+if [[ $(query_nvidia "driver_version") == "" ]]; then
+  echo "Failed to detect GPU driver."
+  exit 1
+fi
+
+
 # Clean up the builder's version of resolv.conf
 sudo rm /mnt/image/etc/resolv.conf
 
 # Skip unmounting:
 #  Sometimes systemd starts, making it hard to unmount
 #  In any case we'll unmount cleanly when the instance shuts down
+
 echo IMAGE_WAS_CREATED
diff --git a/tools/create_base_image_hostlib.sh b/tools/create_base_image_hostlib.sh
index d8e583c7..e5d2e3c5 100755
--- a/tools/create_base_image_hostlib.sh
+++ b/tools/create_base_image_hostlib.sh
@@ -22,12 +22,12 @@ DEFINE_string dest_project "$(gcloud config get-value project)" \
   "Project to use for the new image" "p"
 DEFINE_string launch_instance "" \
   "Name of the instance to launch with the new image" "l"
-DEFINE_string source_image_family debian-9 "Image familty to use as the base" \
-  "s"
+DEFINE_string source_image_family "debian-10" \
+  "Image familty to use as the base" "s"
 DEFINE_string source_image_project debian-cloud \
   "Project holding the base image" "m"
 DEFINE_string repository_url \
-  https://github.com/google/android-cuttlefish.git \
+  "https://github.com/google/android-cuttlefish.git" \
   "URL to the repository with host changes" "u"
 DEFINE_string repository_branch master \
   "Branch to check out" "b"
@@ -102,6 +102,8 @@ main() {
     --image-family="${FLAGS_source_image_family}" \
     --image-project="${FLAGS_source_image_project}" \
     --boot-disk-size=200GiB \
+    --accelerator="type=nvidia-tesla-p100-vws,count=1" \
+    --maintenance-policy=TERMINATE \
     "${FLAGS_build_instance}"
   wait_for_instance "${PZ[@]}" "${FLAGS_build_instance}"
   # Ubuntu tends to mount the wrong disk as root, so help it by waiting until
@@ -133,6 +135,8 @@ main() {
       --image="${FLAGS_dest_image}" \
       --machine-type=n1-standard-4 \
       --scopes storage-ro \
+      --accelerator="type=nvidia-tesla-p100-vws,count=1" \
+      --maintenance-policy=TERMINATE \
       "${FLAGS_launch_instance}"
   fi
   cat <<EOF