[autoscaler][AWS] Make sure subnets belong to same VPC as user-specified security groups (#13558)

* initial commit

* Filter subnets by security groups' VPCs

* fix stubs

* wip

* Fix inbound rule logic. Tests WIP.

* wip

* unit test

* example yaml

* Unit test tests for bug being fixed

* Update python/ray/tests/aws/utils/constants.py

Co-authored-by: Thomas Desrosiers <681004+thomasdesr@users.noreply.github.com>

Co-authored-by: Thomas Desrosiers <681004+thomasdesr@users.noreply.github.com>

python/ray/autoscaler/_private/aws/config.py

@@ -5,6 +5,7 @@ import itertools
 import json
 import os
 import time
+from typing import Any, Dict, List
 import logging
 
 import boto3
@@ -357,9 +358,23 @@ def _configure_subnet(config):
     ec2 = _resource("ec2", config)
     use_internal_ips = config["provider"].get("use_internal_ips", False)
 
+    # If head or worker security group is specified, filter down to subnets
+    # belonging to the same VPC as the security group.
+    sg_ids = (config["head_node"].get("SecurityGroupIds", []) +
+              config["worker_nodes"].get("SecurityGroupIds", []))
+    if sg_ids:
+        vpc_id_of_sg = _get_vpc_id_of_sg(sg_ids, config)
+    else:
+        vpc_id_of_sg = None
+
     try:
+        candidate_subnets = ec2.subnets.all()
+        if vpc_id_of_sg:
+            candidate_subnets = [
+                s for s in candidate_subnets if s.vpc_id == vpc_id_of_sg
+            ]
         subnets = sorted(
-            (s for s in ec2.subnets.all() if s.state == "available" and (
+            (s for s in candidate_subnets if s.state == "available" and (
                 use_internal_ips or s.map_public_ip_on_launch)),
             reverse=True,  # sort from Z-A
             key=lambda subnet: subnet.availability_zone)
@@ -414,6 +429,34 @@ def _configure_subnet(config):
     return config
 
 
+def _get_vpc_id_of_sg(sg_ids: List[str], config: Dict[str, Any]) -> str:
+    """Returns the VPC id of the security groups with the provided security
+    group ids.
+
+    Errors if the provided security groups belong to multiple VPCs.
+    Errors if no security group with any of the provided ids is identified.
+    """
+    sg_ids = list(set(sg_ids))
+
+    ec2 = _resource("ec2", config)
+    filters = [{"Name": "group-id", "Values": sg_ids}]
+    security_groups = ec2.security_groups.filter(Filters=filters)
+    vpc_ids = [sg.vpc_id for sg in security_groups]
+    vpc_ids = list(set(vpc_ids))
+
+    multiple_vpc_msg = "All security groups specified in the cluster config "\
+        "should belong to the same VPC."
+    cli_logger.doassert(len(vpc_ids) <= 1, multiple_vpc_msg)
+    assert len(vpc_ids) <= 1, multiple_vpc_msg
+
+    no_sg_msg = "Failed to detect a security group with id equal to any of "\
+        "the configured SecurityGroupIds."
+    cli_logger.doassert(len(vpc_ids) > 0, no_sg_msg)
+    assert len(vpc_ids) > 0, no_sg_msg
+
+    return vpc_ids[0]
+
+
 def _configure_security_group(config):
     _set_config_info(
         head_security_group_src="config", workers_security_group_src="config")
@@ -566,6 +609,13 @@ def _create_security_group(config, vpc_id, group_name):
 
 def _upsert_security_group_rules(conf, security_groups):
     sgids = {sg.id for sg in security_groups.values()}
+
+    # Update sgids to include user-specified security groups.
+    # This is necessary if the user specifies the head node type's security
+    # groups but not the worker's, or vice-versa.
+    for node_type in NODE_KIND_CONFIG_KEYS.values():
+        sgids.update(conf[node_type].get("SecurityGroupIds", []))
+
     # sort security group items for deterministic inbound rule config order
     # (mainly supports more precise stub-based boto3 unit testing)
     for node_type, sg in sorted(security_groups.items()):
@@ -583,7 +633,7 @@ def _update_inbound_rules(target_security_group, sgids, config):
 
 
 def _create_default_inbound_rules(sgids, extended_rules=[]):
-    intracluster_rules = _create_default_instracluster_inbound_rules(sgids)
+    intracluster_rules = _create_default_intracluster_inbound_rules(sgids)
     ssh_rules = _create_default_ssh_inbound_rules()
     merged_rules = itertools.chain(
         intracluster_rules,
@@ -593,7 +643,7 @@ def _create_default_inbound_rules(sgids, extended_rules=[]):
     return list(merged_rules)
 
 
-def _create_default_instracluster_inbound_rules(intracluster_sgids):
+def _create_default_intracluster_inbound_rules(intracluster_sgids):
     return [{
         "FromPort": -1,
         "ToPort": -1,

python/ray/autoscaler/aws/example-head-and-worker-security-group.yaml

@@ -0,0 +1,31 @@
+cluster_name: sg
+
+max_workers: 1
+
+provider:
+    type: aws
+    region: us-west-2
+    availability_zone: us-west-2a
+
+auth:
+    ssh_user: ubuntu
+
+# If required, head and worker nodes can exist on subnets in different VPCs and
+# communicate via VPC peering.
+
+# VPC peering overview: https://docs.aws.amazon.com/vpc/latest/userguide/vpc-peering.html.
+# Setup VPC peering: https://docs.aws.amazon.com/vpc/latest/peering/create-vpc-peering-connection.html.
+# Configure VPC peering route tables: https://docs.aws.amazon.com/vpc/latest/peering/vpc-peering-routing.html.
+
+# To enable external SSH connectivity, you should also ensure that your VPC
+# is configured to assign public IPv4 addresses to every EC2 instance
+# assigned to it.
+head_node:
+    SecurityGroupIds:
+        - sg-1234abcd # Replace with an actual security group id.
+
+worker_nodes:
+    SecurityGroupIds:
+        - sg-1234abcd # Replace with an actual security group id.
+
+