From cd86578b85d4c3a9ab7df85ae72ada99fd41b15a Mon Sep 17 00:00:00 2001
From: Dmitri Gekhtman <62982571+DmitriGekhtman@users.noreply.github.com>
Date: Wed, 24 Mar 2021 00:16:25 -0400
Subject: [PATCH] [autoscaler][aws] Use subnets in only one VPC (#14868)

---
 python/ray/autoscaler/_private/aws/config.py |  9 ++-
 python/ray/tests/aws/test_autoscaler_aws.py  | 61 +++++++++++++++++++-
 python/ray/tests/aws/utils/stubs.py          | 12 ++++
 3 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/python/ray/autoscaler/_private/aws/config.py b/python/ray/autoscaler/_private/aws/config.py
index cd9d9c1c11cea..83b25d11605ba 100644
--- a/python/ray/autoscaler/_private/aws/config.py
+++ b/python/ray/autoscaler/_private/aws/config.py
@@ -415,7 +415,14 @@ def _configure_subnet(config):
                 "to populate the list of subnets and trying this again.".
                 format(config["provider"]["availability_zone"]))
 
-    subnet_ids = [s.subnet_id for s in subnets]
+    # Use subnets in only one VPC, so that _configure_security_groups only
+    # needs to create a security group in this one VPC. Otherwise, we'd need
+    # to set up security groups in all of the user's VPCs and set up networking
+    # rules to allow traffic between these groups.
+    # See https://github.com/ray-project/ray/pull/14868.
+    subnet_ids = [
+        s.subnet_id for s in subnets if s.vpc_id == subnets[0].vpc_id
+    ]
     if "SubnetIds" not in config["head_node"]:
         _set_config_info(head_subnet_src="default")
         config["head_node"]["SubnetIds"] = subnet_ids
diff --git a/python/ray/tests/aws/test_autoscaler_aws.py b/python/ray/tests/aws/test_autoscaler_aws.py
index 09c736a07c06b..78fd71cd31ee5 100644
--- a/python/ray/tests/aws/test_autoscaler_aws.py
+++ b/python/ray/tests/aws/test_autoscaler_aws.py
@@ -8,7 +8,66 @@
 from ray.tests.aws.utils.constants import AUX_SUBNET, DEFAULT_SUBNET, \
     DEFAULT_SG_AUX_SUBNET, DEFAULT_SG, DEFAULT_SG_DUAL_GROUP_RULES, \
     DEFAULT_SG_WITH_RULES_AUX_SUBNET, AUX_SG, \
-    DEFAULT_SG_WITH_NAME, DEFAULT_SG_WITH_NAME_AND_RULES, CUSTOM_IN_BOUND_RULES
+    DEFAULT_SG_WITH_RULES, DEFAULT_SG_WITH_NAME, \
+    DEFAULT_SG_WITH_NAME_AND_RULES, CUSTOM_IN_BOUND_RULES
+
+
+def test_use_subnets_in_only_one_vpc(iam_client_stub, ec2_client_stub):
+    """
+    This test validates that when bootstrap_aws populates the SubnetIds field,
+    all of the subnets used belong to the same VPC, and that a SecurityGroup
+    in that VPC is correctly configured.
+    """
+    stubs.configure_iam_role_default(iam_client_stub)
+    stubs.configure_key_pair_default(ec2_client_stub)
+
+    # Add a response with a thousand subnets all in different VPCs.
+    # After filtering, only subnet in one particular VPC should remain.
+    # Thus head_node.SubnetIds and worker_nodes.SubnetIds should end up as
+    # being length-one lists after the bootstrap_config.
+    stubs.describe_a_thousand_subnets_in_different_vpcs(ec2_client_stub)
+
+    # describe the subnet in use while determining its vpc
+    stubs.describe_subnets_echo(ec2_client_stub, DEFAULT_SUBNET)
+    # given no existing security groups within the VPC...
+    stubs.describe_no_security_groups(ec2_client_stub)
+    # expect to create a security group on the VPC
+    stubs.create_sg_echo(ec2_client_stub, DEFAULT_SG)
+    # expect new security group details to be retrieved after creation
+    stubs.describe_sgs_on_vpc(
+        ec2_client_stub,
+        [DEFAULT_SUBNET["VpcId"]],
+        [DEFAULT_SG],
+    )
+
+    # given no existing default security group inbound rules...
+    # expect to authorize all default inbound rules
+    stubs.authorize_sg_ingress(
+        ec2_client_stub,
+        DEFAULT_SG_WITH_RULES,
+    )
+
+    # expect another call to describe the above security group while checking
+    # a second time if it has ip_permissions set ("if not sg.ip_permissions")
+    stubs.describe_an_sg_2(
+        ec2_client_stub,
+        DEFAULT_SG_WITH_RULES,
+    )
+
+    # given our mocks and an example config file as input...
+    # expect the config to be loaded, validated, and bootstrapped successfully
+    config = helpers.bootstrap_aws_example_config_file("example-full.yaml")
+    _get_vpc_id_or_die.cache_clear()
+
+    # We've filtered down to only one subnet id -- only one of the thousand
+    # subnets generated by ec2.subnets.all() belongs to the right VPC.
+    assert config["head_node"]["SubnetIds"] == [DEFAULT_SUBNET["SubnetId"]]
+    assert config["worker_nodes"]["SubnetIds"] == [DEFAULT_SUBNET["SubnetId"]]
+    # Check that the security group has been filled correctly.
+    assert config["head_node"]["SecurityGroupIds"] == [DEFAULT_SG["GroupId"]]
+    assert config["worker_nodes"]["SecurityGroupIds"] == [
+        DEFAULT_SG["GroupId"]
+    ]
 
 
 def test_create_sg_different_vpc_same_rules(iam_client_stub, ec2_client_stub):
diff --git a/python/ray/tests/aws/utils/stubs.py b/python/ray/tests/aws/utils/stubs.py
index 61f1f9ab632b7..1c48ed47667b0 100644
--- a/python/ray/tests/aws/utils/stubs.py
+++ b/python/ray/tests/aws/utils/stubs.py
@@ -85,6 +85,18 @@ def describe_a_security_group(ec2_client_stub, security_group):
         service_response={"SecurityGroups": [security_group]})
 
 
+def describe_an_sg_2(ec2_client_stub, security_group):
+    """Same as last function, different input param format.
+
+    A call with this input parameter format is made when sg.ip_permissions is
+    accessed in aws/config.py.
+    """
+    ec2_client_stub.add_response(
+        "describe_security_groups",
+        expected_params={"GroupIds": [security_group["GroupId"]]},
+        service_response={"SecurityGroups": [security_group]})
+
+
 def create_sg_echo(ec2_client_stub, security_group):
     ec2_client_stub.add_response(
         "create_security_group",