Add support for additional disks for config-local (#5030) (#5040)

* Add support for additional disks for config-local

(cherry picked from commit e4861f52e0)

net/scripts/gce-provider.sh

@@ -126,6 +126,7 @@ cloud_CreateInstances() {
   declare optionalStartupScript="$8"
   declare optionalAddress="$9"
   declare optionalBootDiskType="${10}"
+  declare optionalAdditionalDiskSize="${11}"
 
   if $enableGpu; then
     # Custom Ubuntu 18.04 LTS image with CUDA 9.2 and CUDA 10.0 installed
@@ -198,6 +199,22 @@ cloud_CreateInstances() {
     set -x
     gcloud beta compute instances create "${nodes[@]}" "${args[@]}"
   )
+
+  if [[ -n $optionalAdditionalDiskSize ]]; then
+    if [[ $numNodes = 1 ]]; then
+      (
+        set -x
+        cloud_CreateAndAttachPersistentDisk "${namePrefix}" "$optionalAdditionalDiskSize" "pd-ssd" "$zone"
+      )
+    else
+      for node in $(seq -f "${namePrefix}%0${#numNodes}g" 1 "$numNodes"); do
+        (
+          set -x
+          cloud_CreateAndAttachPersistentDisk "${node}" "$optionalAdditionalDiskSize" "pd-ssd" "$zone"
+        )
+      done
+    fi
+  fi
 }
 
 #
@@ -256,3 +273,31 @@ cloud_FetchFile() {
     gcloud compute scp --zone "$zone" "$instanceName:$remoteFile" "$localFile"
   )
 }
+
+#
+# cloud_CreateAndAttachPersistentDisk [instanceName] [diskSize] [diskType]
+#
+# Create a persistent disk and attach it to a pre-existing VM instance.
+# Set disk to auto-delete upon instance deletion
+#
+cloud_CreateAndAttachPersistentDisk() {
+  declare instanceName="$1"
+  declare diskSize="$2"
+  declare diskType="$3"
+  declare zone="$4"
+  diskName="${instanceName}-pd"
+
+  gcloud beta compute disks create "$diskName" \
+    --size "$diskSize" \
+    --type "$diskType" \
+    --zone "$zone"
+
+  gcloud compute instances attach-disk "$instanceName" \
+    --disk "$diskName" \
+    --zone "$zone"
+
+  gcloud compute instances set-disk-auto-delete "$instanceName" \
+    --disk "$diskName" \
+    --zone "$zone" \
+    --auto-delete
+}