From 7e2b65374dc02dd3d9fdab4dcac1e918321edab1 Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Wed, 29 Aug 2018 08:13:35 -0700
Subject: [PATCH] gce instance types are now configurable

---
 net/gce.sh        | 32 ++++++++++++++++++++++++--------
 scripts/gcloud.sh | 23 ++++++++++++++++++-----
 2 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/net/gce.sh b/net/gce.sh
index 601607f43..dd7d13639 100755
--- a/net/gce.sh
+++ b/net/gce.sh
@@ -7,8 +7,14 @@ source "$here"/../scripts/gcloud.sh
 source "$here"/common.sh
 
 prefix=testnet-dev-$(whoami | sed -e s/[^a-z0-9].*//)
-validatorNodeCount=
-clientNodeCount=
+validatorNodeCount=5
+clientNodeCount=1
+leaderMachineType=n1-standard-1
+leaderAccelerator= # "count=4,type=nvidia-tesla-k80"
+validatorMachineType=$leaderMachineType
+validatorAccelerator=
+clientMachineType=$leaderMachineType
+clientAccelerator=
 
 imageName="ubuntu-16-04-cuda-9-2-new"
 internalNetwork=false
@@ -34,9 +40,9 @@ Manage a GCE-based testnet
                     (default: $prefix)
 
  create-specific options:
- -n number        - Number of validator nodes
- -c number        - Number of client nodes
- -P               - Use GCE internal/private network
+ -n number        - Number of validator nodes (default: $validatorNodeCount)
+ -c number        - Number of client nodes (default: $clientNodeCount)
+ -P               - Use GCE internal/private network (default: $internalNetwork)
  -z               - GCP Zone for the nodes (default: $zone)
  -i imageName     - Existing image on GCE (default: $imageName)
 
@@ -156,11 +162,21 @@ delete)
 create)
   [[ -n $validatorNodeCount ]] || usage "Need number of nodes"
 
-  gcloud_CreateInstances "$prefix-leader" 1 "$zone" "$imageName"
-  gcloud_CreateInstances "$prefix-validator" "$validatorNodeCount" "$zone" "$imageName"
+  echo "Network composition:"
+  echo "Leader = $leaderMachineType (GPU=${leaderAccelerator:-none})"
+  echo "Validators = $validatorNodeCount x $validatorMachineType (GPU=${validatorAccelerator:-none})"
+  echo "Client(s) = $clientNodeCount x $clientMachineType (GPU=${clientAccelerator:-none})"
+  echo ==================================================================
+  echo
+  gcloud_CreateInstances "$prefix-leader" 1 \
+    "$zone" "$imageName" "$leaderMachineType" "$leaderAccelerator"
+  gcloud_CreateInstances "$prefix-validator" "$validatorNodeCount" \
+    "$zone" "$imageName" "$validatorMachineType" "$validatorAccelerator"
   if [[ -n $clientNodeCount ]]; then
-    gcloud_CreateInstances "$prefix-client" "$clientNodeCount" "$zone" "$imageName"
+    gcloud_CreateInstances "$prefix-client" "$clientNodeCount" \
+      "$zone" "$imageName" "$clientMachineType" "$clientAccelerator"
   fi
+
   writeConfigFile
   ;;
 
diff --git a/scripts/gcloud.sh b/scripts/gcloud.sh
index 9842e9f50..6b5006713 100644
--- a/scripts/gcloud.sh
+++ b/scripts/gcloud.sh
@@ -75,7 +75,7 @@ gcloud_ForEachInstance() {
 }
 
 #
-# gcloud_CreateInstances [namePrefix] [numNodes] [zone] [imageName]
+# gcloud_CreateInstances [namePrefix] [numNodes] [zone] [imageName] [machineType] [accelerator]
 #
 # Creates one more identical instances.
 #
@@ -83,6 +83,9 @@ gcloud_ForEachInstance() {
 # numNodes     - number of instances to create
 # zone         - zone to create the instances in
 # imageName    - Disk image for the instances
+# machineType  - GCE machine type
+# accelerator  - Optional accelerator to attach to the instance(s), see
+#                eg, request 4 K80 GPUs with "count=4,type=nvidia-tesla-k80"
 #
 # Tip: use gcloud_FindInstances to locate the instances once this function
 #      returns
@@ -91,6 +94,8 @@ gcloud_CreateInstances() {
   declare numNodes="$2"
   declare zone="$3"
   declare imageName="$4"
+  declare machineType="$5"
+  declare optionalAccelerator="$6"
 
   declare nodes
   if [[ $numNodes = 1 ]]; then
@@ -99,12 +104,20 @@ gcloud_CreateInstances() {
     read -ra nodes <<<$(seq -f "${namePrefix}%g" 1 "$numNodes")
   fi
 
+  declare -a args
+  args=(
+    "--zone=$zone"
+    "--tags=testnet"
+    "--image=$imageName"
+    "--machine-type=$machineType"
+  )
+  if [[ -n $optionalAccelerator ]]; then
+    args+=("--accelerator=$optionalAccelerator")
+  fi
+
   (
     set -x
-    gcloud beta compute instances create "${nodes[@]}" \
-      --zone="$zone" \
-      --tags=testnet \
-      --image="$imageName"
+    gcloud beta compute instances create "${nodes[@]}" "${args[@]}" \
   )
 }