[Autoscaler] Allow users to set the names for security groups created by ray (#11405)

This commit is contained in:
Allen
2020-10-22 12:28:59 -07:00
committed by GitHub
parent 7111a424af
commit cf2ee94e0c
6 changed files with 170 additions and 6 deletions
+3 -1
View File
@@ -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 = {
@@ -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
+15
View File
@@ -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",
+50 -1
View File
@@ -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__]))
+31 -4
View File
@@ -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
+42
View File
@@ -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