mirror of
https://github.com/wassname/ray.git
synced 2026-06-28 13:02:16 +08:00
[Autoscaler] Allow users to set the names for security groups created by ray (#11405)
This commit is contained in:
@@ -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
|
||||
@@ -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",
|
||||
|
||||
@@ -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__]))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user