[autoscaler] Honor separate head and worker node subnet IDs (#8374)

2026-07-13 11:12:12 +08:00 · 2020-05-28 18:16:46 -07:00
parent ad1c05c8b0
commit 76450c8d47
13 changed files with 594 additions and 69 deletions
@@ -190,6 +190,14 @@ py_test(
    deps = ["//:ray_lib"],
 )

+
+py_test(
+    name = "test_autoscaler_aws",
+    size = "small",
+    srcs = ["aws/test_autoscaler_aws.py"],
+    deps = ["//:ray_lib"],
+)
+
 py_test(
    name = "test_autoscaler_yaml",
    size = "small",
@@ -0,0 +1,21 @@
+import pytest
+
+from ray.autoscaler.aws.config import _resource_cache
+
+from botocore.stub import Stubber
+
+
+@pytest.fixture()
+def iam_client_stub():
+    resource = _resource_cache("iam", "us-west-2")
+    with Stubber(resource.meta.client) as stubber:
+        yield stubber
+        stubber.assert_no_pending_responses()
+
+
+@pytest.fixture()
+def ec2_client_stub():
+    resource = _resource_cache("ec2", "us-west-2")
+    with Stubber(resource.meta.client) as stubber:
+        yield stubber
+        stubber.assert_no_pending_responses()
@@ -0,0 +1,76 @@
+import pytest
+
+import ray.tests.aws.utils.stubs as stubs
+import ray.tests.aws.utils.helpers as helpers
+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, DEFAULT_SG_WITH_RULES, AUX_SG
+
+
+def test_create_sg_different_vpc_same_rules(iam_client_stub, ec2_client_stub):
+    # use default stubs to skip ahead to security group configuration
+    stubs.skip_to_configure_sg(ec2_client_stub, iam_client_stub)
+
+    # given head and worker nodes with custom subnets defined...
+    # expect to first describe the worker subnet ID
+    stubs.describe_subnets_echo(ec2_client_stub, AUX_SUBNET)
+    # expect to second describe the head subnet ID
+    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 first create a security group on the worker node VPC
+    stubs.create_sg_echo(ec2_client_stub, DEFAULT_SG_AUX_SUBNET)
+    # expect new worker security group details to be retrieved after creation
+    stubs.describe_sgs_on_vpc(
+        ec2_client_stub,
+        [AUX_SUBNET["VpcId"]],
+        [DEFAULT_SG_AUX_SUBNET],
+    )
+    # expect to second create a security group on the head node VPC
+    stubs.create_sg_echo(ec2_client_stub, DEFAULT_SG)
+    # expect new head 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 head security group inbound rules...
+    # expect to authorize all default head inbound rules
+    stubs.authorize_sg_ingress(
+        ec2_client_stub,
+        DEFAULT_SG_DUAL_GROUP_RULES,
+    )
+    # given no existing default worker security group inbound rules...
+    # expect to authorize all default worker inbound rules
+    stubs.authorize_sg_ingress(
+        ec2_client_stub,
+        DEFAULT_SG_WITH_RULES_AUX_SUBNET,
+    )
+
+    # given the prior modification to the head security group...
+    # expect the next read of a head security group property to reload it
+    stubs.describe_sg_echo(ec2_client_stub, DEFAULT_SG_WITH_RULES)
+    # given the prior modification to the worker security group...
+    # expect the next read of a worker security group property to reload it
+    stubs.describe_sg_echo(ec2_client_stub, DEFAULT_SG_WITH_RULES_AUX_SUBNET)
+
+    # 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-subnets.yaml")
+
+    # expect the bootstrapped config to show different head and worker security
+    # groups residing on different subnets
+    assert config["head_node"]["SecurityGroupIds"] == [DEFAULT_SG["GroupId"]]
+    assert config["head_node"]["SubnetIds"] == [DEFAULT_SUBNET["SubnetId"]]
+    assert config["worker_nodes"]["SecurityGroupIds"] == [AUX_SG["GroupId"]]
+    assert config["worker_nodes"]["SubnetIds"] == [AUX_SUBNET["SubnetId"]]
+
+    # expect no pending responses left in IAM or EC2 client stub queues
+    iam_client_stub.assert_no_pending_responses()
+    ec2_client_stub.assert_no_pending_responses()
+
+
+if __name__ == "__main__":
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))
@@ -0,0 +1,129 @@
+import copy
+import ray
+from datetime import datetime
+
+# Override global constants used in AWS autoscaler config artifact names.
+# This helps ensure that any unmocked test doesn't alter non-test artifacts.
+ray.autoscaler.aws.config.RAY = \
+    "ray-autoscaler-aws-test"
+ray.autoscaler.aws.config.DEFAULT_RAY_INSTANCE_PROFILE = \
+    ray.autoscaler.aws.config.RAY + "-v1"
+ray.autoscaler.aws.config.DEFAULT_RAY_IAM_ROLE = \
+    ray.autoscaler.aws.config.RAY + "-v1"
+ray.autoscaler.aws.config.SECURITY_GROUP_TEMPLATE = \
+    ray.autoscaler.aws.config.RAY + "-{}"
+
+# Default IAM instance profile to expose to tests.
+DEFAULT_INSTANCE_PROFILE = {
+    "Arn": "arn:aws:iam::336924118301:instance-profile/ExampleInstanceProfile",
+    "CreateDate": datetime(2013, 6, 12, 23, 52, 2, 2),
+    "InstanceProfileId": "AIPA0000000000EXAMPLE",
+    "InstanceProfileName": "ExampleInstanceProfile",
+    "Path": "/",
+    "Roles": [
+        {
+            "Arn": "arn:aws:iam::123456789012:role/Test-Role",
+            "AssumeRolePolicyDocument": "ExampleAssumeRolePolicyDocument",
+            "CreateDate": datetime(2013, 1, 9, 6, 33, 26, 2),
+            "Path": "/",
+            "RoleId": "AROA0000000000EXAMPLE",
+            "RoleName": "Test-Role",
+        },
+    ]
+}
+
+# Default EC2 key pair to expose to tests.
+DEFAULT_KEY_PAIR = {
+    "KeyFingerprint": "00:11:22:33:44:55:66:77:88:99:AA:BB:CC:DD:EE:FF:00",
+    "KeyName": ray.autoscaler.aws.config.RAY + "_us-west-2",
+}
+
+# Primary EC2 subnet to expose to tests.
+DEFAULT_SUBNET = {
+    "AvailabilityZone": "us-west-2a",
+    "AvailableIpAddressCount": 251,
+    "CidrBlock": "10.0.1.0/24",
+    "DefaultForAz": False,
+    "MapPublicIpOnLaunch": True,
+    "State": "available",
+    "SubnetId": "subnet-0000000",
+    "VpcId": "vpc-0000000",
+}
+
+# Secondary EC2 subnet to expose to tests as required.
+AUX_SUBNET = {
+    "AvailabilityZone": "us-west-2a",
+    "AvailableIpAddressCount": 251,
+    "CidrBlock": "192.168.1.0/24",
+    "DefaultForAz": False,
+    "MapPublicIpOnLaunch": True,
+    "State": "available",
+    "SubnetId": "subnet-fffffff",
+    "VpcId": "vpc-fffffff",
+}
+
+# Default cluster name to expose to tests.
+DEFAULT_CLUSTER_NAME = "test-cluster-name"
+
+# Default security group settings immediately after creation
+# (prior to inbound rule configuration).
+DEFAULT_SG = {
+    "Description": "Auto-created security group for Ray workers",
+    "GroupName": ray.autoscaler.aws.config.RAY + "-" + DEFAULT_CLUSTER_NAME,
+    "OwnerId": "test-owner",
+    "GroupId": "sg-1234abcd",
+    "VpcId": DEFAULT_SUBNET["VpcId"],
+    "IpPermissions": [],
+    "IpPermissionsEgress": [{
+        "FromPort": -1,
+        "ToPort": -1,
+        "IpProtocol": "-1",
+        "IpRanges": [{
+            "CidrIp": "0.0.0.0/0"
+        }]
+    }],
+    "Tags": []
+}
+
+# Secondary security group settings after creation
+# (prior to inbound rule configuration).
+AUX_SG = copy.deepcopy(DEFAULT_SG)
+AUX_SG["GroupName"] += "-aux"
+AUX_SG["GroupId"] = "sg-dcba4321"
+
+# Default security group settings immediately after creation on aux subnet
+# (prior to inbound rule configuration).
+DEFAULT_SG_AUX_SUBNET = copy.deepcopy(DEFAULT_SG)
+DEFAULT_SG_AUX_SUBNET["VpcId"] = AUX_SUBNET["VpcId"]
+DEFAULT_SG_AUX_SUBNET["GroupId"] = AUX_SG["GroupId"]
+
+# Default security group settings once default inbound rules are applied
+# (if used by both head and worker nodes)
+DEFAULT_SG_WITH_RULES = copy.deepcopy(DEFAULT_SG)
+DEFAULT_SG_WITH_RULES["IpPermissions"] = [{
+    "FromPort": -1,
+    "ToPort": -1,
+    "IpProtocol": "-1",
+    "UserIdGroupPairs": [{
+        "GroupId": DEFAULT_SG["GroupId"]
+    }]
+}, {
+    "FromPort": 22,
+    "ToPort": 22,
+    "IpProtocol": "tcp",
+    "IpRanges": [{
+        "CidrIp": "0.0.0.0/0"
+    }]
+}]
+
+# Default security group once default inbound rules are applied
+# (if using separate security groups for head and worker nodes).
+DEFAULT_SG_DUAL_GROUP_RULES = copy.deepcopy(DEFAULT_SG_WITH_RULES)
+DEFAULT_SG_DUAL_GROUP_RULES["IpPermissions"][0]["UserIdGroupPairs"].append({
+    "GroupId": AUX_SG["GroupId"]
+})
+
+# Default security group on aux subnet once default inbound rules are applied.
+DEFAULT_SG_WITH_RULES_AUX_SUBNET = copy.deepcopy(DEFAULT_SG_DUAL_GROUP_RULES)
+DEFAULT_SG_WITH_RULES_AUX_SUBNET["VpcId"] = AUX_SUBNET["VpcId"]
+DEFAULT_SG_WITH_RULES_AUX_SUBNET["GroupId"] = AUX_SG["GroupId"]
@@ -0,0 +1,28 @@
+import os
+import yaml
+import ray
+
+from ray.autoscaler.commands import fillout_defaults, validate_config
+from ray.tests.aws.utils.constants import DEFAULT_CLUSTER_NAME
+
+
+def get_aws_example_config_file_path(file_name):
+    return os.path.join(
+        os.path.dirname(ray.autoscaler.aws.__file__), file_name)
+
+
+def load_aws_example_config_file(file_name):
+    config_file_path = get_aws_example_config_file_path(file_name)
+    return yaml.safe_load(open(config_file_path).read())
+
+
+def bootstrap_aws_config(config):
+    config = fillout_defaults(config)
+    validate_config(config)
+    config["cluster_name"] = DEFAULT_CLUSTER_NAME
+    return ray.autoscaler.aws.config.bootstrap_aws(config)
+
+
+def bootstrap_aws_example_config_file(file_name):
+    config = load_aws_example_config_file(file_name)
+    return bootstrap_aws_config(config)
@@ -0,0 +1,7 @@
+from ray.autoscaler.aws.config import key_pair
+from ray.tests.aws.utils.constants import DEFAULT_KEY_PAIR
+
+
+def mock_path_exists_key_pair(path):
+    key_name, key_path = key_pair(0, "us-west-2", DEFAULT_KEY_PAIR["KeyName"])
+    return path == key_path
@@ -0,0 +1,104 @@
+import ray
+from ray.tests.aws.utils.mocks import mock_path_exists_key_pair
+from ray.tests.aws.utils.constants import DEFAULT_INSTANCE_PROFILE, \
+    DEFAULT_KEY_PAIR, DEFAULT_SUBNET
+
+from unittest import mock
+
+from botocore.stub import ANY
+
+
+def configure_iam_role_default(iam_client_stub):
+    iam_client_stub.add_response(
+        "get_instance_profile",
+        expected_params={
+            "InstanceProfileName": ray.autoscaler.aws.config.
+            DEFAULT_RAY_INSTANCE_PROFILE
+        },
+        service_response={"InstanceProfile": DEFAULT_INSTANCE_PROFILE})
+
+
+def configure_key_pair_default(ec2_client_stub):
+    patcher = mock.patch("os.path.exists")
+    os_path_exists_mock = patcher.start()
+    os_path_exists_mock.side_effect = mock_path_exists_key_pair
+
+    ec2_client_stub.add_response(
+        "describe_key_pairs",
+        expected_params={
+            "Filters": [{
+                "Name": "key-name",
+                "Values": [DEFAULT_KEY_PAIR["KeyName"]]
+            }]
+        },
+        service_response={"KeyPairs": [DEFAULT_KEY_PAIR]})
+
+
+def configure_subnet_default(ec2_client_stub):
+    ec2_client_stub.add_response(
+        "describe_subnets",
+        expected_params={},
+        service_response={"Subnets": [DEFAULT_SUBNET]})
+
+
+def skip_to_configure_sg(ec2_client_stub, iam_client_stub):
+    configure_iam_role_default(iam_client_stub)
+    configure_key_pair_default(ec2_client_stub)
+    configure_subnet_default(ec2_client_stub)
+
+
+def describe_subnets_echo(ec2_client_stub, subnet):
+    ec2_client_stub.add_response(
+        "describe_subnets",
+        expected_params={
+            "Filters": [{
+                "Name": "subnet-id",
+                "Values": [subnet["SubnetId"]]
+            }]
+        },
+        service_response={"Subnets": [subnet]})
+
+
+def describe_no_security_groups(ec2_client_stub):
+    ec2_client_stub.add_response(
+        "describe_security_groups",
+        expected_params={"Filters": ANY},
+        service_response={})
+
+
+def create_sg_echo(ec2_client_stub, security_group):
+    ec2_client_stub.add_response(
+        "create_security_group",
+        expected_params={
+            "Description": security_group["Description"],
+            "GroupName": security_group["GroupName"],
+            "VpcId": security_group["VpcId"]
+        },
+        service_response={"GroupId": security_group["GroupId"]})
+
+
+def describe_sgs_on_vpc(ec2_client_stub, vpc_ids, security_groups):
+    ec2_client_stub.add_response(
+        "describe_security_groups",
+        expected_params={"Filters": [{
+            "Name": "vpc-id",
+            "Values": vpc_ids
+        }]},
+        service_response={"SecurityGroups": security_groups})
+
+
+def authorize_sg_ingress(ec2_client_stub, security_group):
+    ec2_client_stub.add_response(
+        "authorize_security_group_ingress",
+        expected_params={
+            "GroupId": security_group["GroupId"],
+            "IpPermissions": security_group["IpPermissions"]
+        },
+        service_response={})
+
+
+def describe_sg_echo(ec2_client_stub, security_group):
+    ec2_client_stub.add_response(
+        "describe_security_groups",
+        expected_params={"GroupIds": [security_group["GroupId"]]},
+        service_response={"SecurityGroups": [security_group]})