diff --git a/python/ray/autoscaler/_private/aws/config.py b/python/ray/autoscaler/_private/aws/config.py index efde00fa6..89ebca7f1 100644 --- a/python/ray/autoscaler/_private/aws/config.py +++ b/python/ray/autoscaler/_private/aws/config.py @@ -523,7 +523,9 @@ def _get_or_create_vpc_security_groups(conf, node_types): } # Generate the name of the security group we're looking for... - expected_sg_name = SECURITY_GROUP_TEMPLATE.format(conf["cluster_name"]) + expected_sg_name = conf["provider"] \ + .get("security_group", {}) \ + .get("GroupName", SECURITY_GROUP_TEMPLATE.format(conf["cluster_name"])) # Figure out which security groups with this name exist for each VPC... vpc_to_existing_sg = { diff --git a/python/ray/autoscaler/aws/example-security-group.yaml b/python/ray/autoscaler/aws/example-security-group.yaml new file mode 100644 index 000000000..528ab99f3 --- /dev/null +++ b/python/ray/autoscaler/aws/example-security-group.yaml @@ -0,0 +1,29 @@ +# An unique identifier for the head node and workers of this cluster. +cluster_name: minimal + +# The maximum number of workers nodes to launch in addition to the head +# node. This takes precedence over min_workers. min_workers default to 0. +max_workers: 1 + +# Cloud-provider specific configuration. +provider: + type: aws + region: us-west-2 + availability_zone: us-west-2a + # Security group to create with custom in bound rules and name. + security_group: + GroupName: test_security_group_name + IpPermissions: + - FromPort: 443 + ToPort: 443 + IpProtocol: TCP + IpRanges: + - CidrIp: 0.0.0.0/0 + - FromPort: 8265 + ToPort: 8265 + IpProtocol: TCP + IpRanges: + - CidrIp: 0.0.0.0/0 +# How Ray will authenticate with newly launched nodes. +auth: + ssh_user: ubuntu diff --git a/python/ray/autoscaler/ray-schema.json b/python/ray/autoscaler/ray-schema.json index 640074ffd..276e4c654 100644 --- a/python/ray/autoscaler/ray-schema.json +++ b/python/ray/autoscaler/ray-schema.json @@ -140,6 +140,21 @@ "type": ["string", "null"], "description": "GCP globally unique project id" }, + "security_group": { + "type": "object", + "description": "AWS security group", + "additionalProperties": false, + "properties": { + "GroupName": { + "type": "string", + "description": "Security group name" + }, + "IpPermissions": { + "type": "array", + "description": "Security group in bound rules" + } + } + }, "gcp_credentials": { "type": "object", "description": "Credentials for authenticating with the GCP client", diff --git a/python/ray/tests/aws/test_autoscaler_aws.py b/python/ray/tests/aws/test_autoscaler_aws.py index b8ad6f31e..992b90f09 100644 --- a/python/ray/tests/aws/test_autoscaler_aws.py +++ b/python/ray/tests/aws/test_autoscaler_aws.py @@ -1,10 +1,12 @@ import pytest +from ray.autoscaler._private.aws.config import _get_vpc_id_or_die 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 + DEFAULT_SG_WITH_RULES_AUX_SUBNET, DEFAULT_SG_WITH_RULES, AUX_SG, \ + DEFAULT_SG_WITH_NAME, DEFAULT_SG_WITH_NAME_AND_RULES, CUSTOM_IN_BOUND_RULES def test_create_sg_different_vpc_same_rules(iam_client_stub, ec2_client_stub): @@ -71,6 +73,53 @@ def test_create_sg_different_vpc_same_rules(iam_client_stub, ec2_client_stub): ec2_client_stub.assert_no_pending_responses() +def test_create_sg_with_custom_inbound_rules_and_name(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) + + # expect to 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 create a security group on the head node VPC + stubs.create_sg_echo(ec2_client_stub, DEFAULT_SG_WITH_NAME) + # expect new head security group details to be retrieved after creation + stubs.describe_sgs_on_vpc( + ec2_client_stub, + [DEFAULT_SUBNET["VpcId"]], + [DEFAULT_SG_WITH_NAME], + ) + + # given custom existing default head security group inbound rules... + # expect to authorize both default and custom inbound rules + stubs.authorize_sg_ingress( + ec2_client_stub, + DEFAULT_SG_WITH_NAME_AND_RULES, + ) + + # 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_NAME_AND_RULES) + + _get_vpc_id_or_die.cache_clear() + # 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-security-group.yaml") + + # expect the bootstrapped config to have the custom security group... + # name and in bound rules + assert config["provider"]["security_group"][ + "GroupName"] == DEFAULT_SG_WITH_NAME_AND_RULES["GroupName"] + assert config["provider"]["security_group"][ + "IpPermissions"] == CUSTOM_IN_BOUND_RULES + + # 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__])) diff --git a/python/ray/tests/aws/utils/constants.py b/python/ray/tests/aws/utils/constants.py index 0267b8518..cdcf5a79c 100644 --- a/python/ray/tests/aws/utils/constants.py +++ b/python/ray/tests/aws/utils/constants.py @@ -98,10 +98,7 @@ 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"] = [{ +DEFAULT_IN_BOUND_RULES = [{ "FromPort": -1, "ToPort": -1, "IpProtocol": "-1", @@ -116,6 +113,10 @@ DEFAULT_SG_WITH_RULES["IpPermissions"] = [{ "CidrIp": "0.0.0.0/0" }] }] +# 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"] = DEFAULT_IN_BOUND_RULES # Default security group once default inbound rules are applied # (if using separate security groups for head and worker nodes). @@ -128,3 +129,29 @@ DEFAULT_SG_DUAL_GROUP_RULES["IpPermissions"][0]["UserIdGroupPairs"].append({ 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"] + +# Default security group with custom name +DEFAULT_SG_WITH_NAME = copy.deepcopy(DEFAULT_SG) +DEFAULT_SG_WITH_NAME["GroupName"] = "test_security_group_name" + +CUSTOM_IN_BOUND_RULES = [{ + "FromPort": 443, + "ToPort": 443, + "IpProtocol": "TCP", + "IpRanges": [{ + "CidrIp": "0.0.0.0/0" + }] +}, { + "FromPort": 8265, + "ToPort": 8265, + "IpProtocol": "TCP", + "IpRanges": [{ + "CidrIp": "0.0.0.0/0" + }] +}] + +# Default security group with custom name once... +# default and custom in bound rules are applied +DEFAULT_SG_WITH_NAME_AND_RULES = copy.deepcopy(DEFAULT_SG_WITH_NAME) +DEFAULT_SG_WITH_NAME_AND_RULES[ + "IpPermissions"] = DEFAULT_IN_BOUND_RULES + CUSTOM_IN_BOUND_RULES diff --git a/python/ray/tests/test_autoscaler_yaml.py b/python/ray/tests/test_autoscaler_yaml.py index 5b9b99ee0..8d31fc076 100644 --- a/python/ray/tests/test_autoscaler_yaml.py +++ b/python/ray/tests/test_autoscaler_yaml.py @@ -138,6 +138,48 @@ class AutoscalingConfigTest(unittest.TestCase): self._test_invalid_config( os.path.join("tests", "additional_property.yaml")) + @unittest.skipIf(sys.platform == "win32", "Failing on Windows.") + def testValidateCustomSecurityGroupConfig(self): + aws_config_path = os.path.join(RAY_PATH, + "autoscaler/aws/example-minimal.yaml") + with open(aws_config_path) as f: + config = yaml.safe_load(f) + + # Test validate security group with custom permissions + ip_permissions = [{ + "FromPort": port, + "ToPort": port, + "IpProtocol": "TCP", + "IpRanges": [{ + "CidrIp": "0.0.0.0/0" + }], + } for port in [80, 443, 8265]] + config["provider"].update({ + "security_group": { + "IpPermissions": ip_permissions + } + }) + config = prepare_config(copy.deepcopy(config)) + try: + validate_config(config) + assert config["provider"]["security_group"][ + "IpPermissions"] == ip_permissions + except Exception: + self.fail( + "Failed to validate config with security group in bound rules!" + ) + + # Test validate security group with custom name + group_name = "test_security_group_name" + config["provider"]["security_group"].update({"GroupName": group_name}) + + try: + validate_config(config) + assert config["provider"]["security_group"][ + "GroupName"] == group_name + except Exception: + self.fail("Failed to validate config with security group name!") + if __name__ == "__main__": import pytest