From ec38dba20913aa1dab5ef8ccc68562f520a8c877 Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Wed, 5 Sep 2018 09:31:50 -0700
Subject: [PATCH] GCE leader nodes can now be provisioned with a static IP
 address

---
 net/gce.sh        | 31 ++++++++++++++++++++-----------
 scripts/gcloud.sh | 17 ++++++++++++++++-
 2 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/net/gce.sh b/net/gce.sh
index 791958459..862f2a94b 100755
--- a/net/gce.sh
+++ b/net/gce.sh
@@ -19,6 +19,7 @@ clientAccelerator=
 imageName="ubuntu-16-04-cuda-9-2-new"
 publicNetwork=false
 zone="us-west1-b"
+leaderAddress=
 
 usage() {
   exitcode=0
@@ -40,12 +41,13 @@ Configure a GCE-based testnet
                       (default: $prefix)
 
  create-specific options:
-   -n number        - Number of validator nodes (default: $validatorNodeCount)
-   -c number        - Number of client nodes (default: $clientNodeCount)
+   -n [number]      - Number of validator nodes (default: $validatorNodeCount)
+   -c [number]      - Number of client nodes (default: $clientNodeCount)
    -P               - Use public network IP addresses (default: $publicNetwork)
-   -z               - GCP Zone for the nodes (default: $zone)
-   -i imageName     - Existing image on GCE (default: $imageName)
+   -z [zone]        - GCP Zone for the nodes (default: $zone)
+   -i [imageName]   - Existing image on GCE (default: $imageName)
    -g               - Enable GPU
+   -a [address]     - Set the leader node's external IP address to this GCE address
 
  config-specific options:
    none
@@ -63,7 +65,7 @@ command=$1
 shift
 [[ $command = create || $command = config || $command = delete ]] || usage "Invalid command: $command"
 
-while getopts "h?p:Pi:n:c:z:g" opt; do
+while getopts "h?p:Pi:n:c:z:ga:" opt; do
   case $opt in
   h | \?)
     usage
@@ -89,6 +91,9 @@ while getopts "h?p:Pi:n:c:z:g" opt; do
   g)
     leaderAccelerator="count=4,type=nvidia-tesla-k80"
     ;;
+  a)
+    leaderAddress=$OPTARG
+    ;;
   *)
     usage "Error: unhandled option: $opt"
     ;;
@@ -187,13 +192,17 @@ create)
   echo "Client(s) = $clientNodeCount x $clientMachineType (GPU=${clientAccelerator:-none})"
   echo ==================================================================
   echo
-  gcloud_CreateInstances "$prefix-leader" 1 \
-    "$zone" "$imageName" "$leaderMachineType" "$leaderAccelerator" "$here/remote/remote-startup.sh"
-  gcloud_CreateInstances "$prefix-validator" "$validatorNodeCount" \
-    "$zone" "$imageName" "$validatorMachineType" "$validatorAccelerator" "$here/remote/remote-startup.sh"
+  gcloud_CreateInstances "$prefix-leader" 1 "$zone" \
+    "$imageName" "$leaderMachineType" "$leaderAccelerator" \
+    "$here/remote/remote-startup.sh" "$leaderAddress" \
+
+  gcloud_CreateInstances "$prefix-validator" "$validatorNodeCount" "$zone" \
+    "$imageName" "$validatorMachineType" "$validatorAccelerator" \
+    "$here/remote/remote-startup.sh" ""
   if [[ -n $clientNodeCount ]]; then
-    gcloud_CreateInstances "$prefix-client" "$clientNodeCount" \
-      "$zone" "$imageName" "$clientMachineType" "$clientAccelerator" "$here/remote/remote-startup.sh"
+    gcloud_CreateInstances "$prefix-client" "$clientNodeCount" "$zone" \
+      "$imageName" "$clientMachineType" "$clientAccelerator" \
+      "$here/remote/remote-startup.sh" ""
   fi
 
   prepareInstancesAndWriteConfigFile
diff --git a/scripts/gcloud.sh b/scripts/gcloud.sh
index be7bc7049..db3aa58c4 100644
--- a/scripts/gcloud.sh
+++ b/scripts/gcloud.sh
@@ -75,7 +75,8 @@ gcloud_ForEachInstance() {
 }
 
 #
-# gcloud_CreateInstances [namePrefix] [numNodes] [zone] [imageName] [machineType] [accelerator] [startupScript]
+# gcloud_CreateInstances [namePrefix] [numNodes] [zone] [imageName]
+#                        [machineType] [accelerator] [startupScript] [address]
 #
 # Creates one more identical instances.
 #
@@ -87,6 +88,9 @@ gcloud_ForEachInstance() {
 # accelerator   - Optional accelerator to attach to the instance(s), see
 #                 eg, request 4 K80 GPUs with "count=4,type=nvidia-tesla-k80"
 # startupScript - Optional startup script to execute when the instance boots
+# address       - Optional name of the GCE static IP address to attach to the
+#                 instance.  Requires that |numNodes| = 1 and that addressName
+#                 has been provisioned in the GCE region that is hosting |zone|
 #
 # Tip: use gcloud_FindInstances to locate the instances once this function
 #      returns
@@ -98,6 +102,7 @@ gcloud_CreateInstances() {
   declare machineType="$5"
   declare optionalAccelerator="$6"
   declare optionalStartupScript="$7"
+  declare optionalAddress="$8"
 
   declare nodes
   if [[ $numNodes = 1 ]]; then
@@ -127,6 +132,16 @@ gcloud_CreateInstances() {
     )
   fi
 
+  if [[ -n $optionalAddress ]]; then
+    [[ $numNodes = 1 ]] || {
+      echo "Error: address may not be supplied when provisioning multiple nodes: $optionalAddress"
+      exit 1
+    }
+    args+=(
+      "--address=$optionalAddress"
+    )
+  fi
+
   (
     set -x
     gcloud beta compute instances create "${nodes[@]}" "${args[@]}"