Added option for availability zone

python/ray/autoscaler/autoscaler.py

@@ -49,6 +49,7 @@
     "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.

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

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: