From a59a9e20aff87a5e47d49e2493a40f74b60a1cdb Mon Sep 17 00:00:00 2001
From: Peter Schafhalter <pschafhalter@berkeley.edu>
Date: Tue, 9 Jan 2018 13:49:47 -0800
Subject: [PATCH] Added option for availability zone (#1393)

---
 python/ray/autoscaler/autoscaler.py    |  1 +
 python/ray/autoscaler/aws/config.py    | 15 ++++++++++++++-
 python/ray/autoscaler/aws/example.yaml |  1 +
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/python/ray/autoscaler/autoscaler.py b/python/ray/autoscaler/autoscaler.py
index 95e4b0bf2..d47857982 100644
--- a/python/ray/autoscaler/autoscaler.py
+++ b/python/ray/autoscaler/autoscaler.py
@@ -49,6 +49,7 @@ CLUSTER_CONFIG_SCHEMA = {
     "provider": {
         "type": str,  # e.g. aws
         "region": str,  # e.g. us-east-1
+        "availability_zone": str,  # e.g. us-east-1a
     },
 
     # How Ray will authenticate with newly launched nodes.
diff --git a/python/ray/autoscaler/aws/config.py b/python/ray/autoscaler/aws/config.py
index b303c101e..765bc8570 100644
--- a/python/ray/autoscaler/aws/config.py
+++ b/python/ray/autoscaler/aws/config.py
@@ -143,7 +143,20 @@ def _configure_subnet(config):
             "No subnets found, try manually creating an instance in "
             "your specified region to populate the list of subnets "
             "and trying this again.")
-    default_subnet = subnets[0]
+    if "availability_zone" in config["provider"]:
+        default_subnet = next((s for s in subnets
+                               if s.availability_zone ==
+                               config["provider"]["availability_zone"]),
+                              None)
+        if not default_subnet:
+            raise Exception(
+                "No available subnets matching availability zone {} "
+                "found. Choose a different availability zone or try "
+                "manually creating an instance in your specified region "
+                "to populate the list of subnets and trying this again."
+                .format(config["provider"]["availability_zone"]))
+    else:
+        default_subnet = subnets[0]
 
     if "SubnetId" not in config["head_node"]:
         config["head_node"]["SubnetId"] = default_subnet.id
diff --git a/python/ray/autoscaler/aws/example.yaml b/python/ray/autoscaler/aws/example.yaml
index 82375664d..43ee464e2 100644
--- a/python/ray/autoscaler/aws/example.yaml
+++ b/python/ray/autoscaler/aws/example.yaml
@@ -22,6 +22,7 @@ idle_timeout_minutes: 5
 provider:
     type: aws
     region: us-west-2
+    # availability_zone: us-west-2a
 
 # How Ray will authenticate with newly launched nodes.
 auth: