wehub-resource-sync
229 lines
7.1 KiB
Python
229 lines
7.1 KiB
Python
import copy
|
|
from datetime import datetime
|
|
|
|
import ray
|
|
from ray.autoscaler.tags import (
|
|
NODE_KIND_HEAD,
|
|
TAG_RAY_LAUNCH_CONFIG,
|
|
TAG_RAY_NODE_KIND,
|
|
TAG_RAY_USER_NODE_TYPE,
|
|
)
|
|
|
|
# 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._private.aws.config.RAY = "ray-autoscaler-aws-test"
|
|
ray.autoscaler._private.aws.config.DEFAULT_RAY_INSTANCE_PROFILE = (
|
|
ray.autoscaler._private.aws.config.RAY + "-v1"
|
|
)
|
|
ray.autoscaler._private.aws.config.DEFAULT_RAY_IAM_ROLE = (
|
|
ray.autoscaler._private.aws.config.RAY + "-v1"
|
|
)
|
|
ray.autoscaler._private.aws.config.SECURITY_GROUP_TEMPLATE = (
|
|
ray.autoscaler._private.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._private.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",
|
|
}
|
|
|
|
|
|
def subnet_in_vpc(vpc_num):
|
|
"""Returns a copy of DEFAULT_SUBNET whose VpcId ends with the digits
|
|
of vpc_num."""
|
|
subnet = copy.copy(DEFAULT_SUBNET)
|
|
subnet["VpcId"] = f"vpc-{vpc_num:07d}"
|
|
return subnet
|
|
|
|
|
|
A_THOUSAND_SUBNETS_IN_DIFFERENT_VPCS = [
|
|
subnet_in_vpc(vpc_num) for vpc_num in range(1, 1000)
|
|
] + [DEFAULT_SUBNET]
|
|
|
|
|
|
def subnet_in_az(idx):
|
|
azs = ["a", "b", "c", "d"]
|
|
subnet = copy.copy(DEFAULT_SUBNET)
|
|
subnet["AvailabilityZone"] = "us-west-2" + azs[idx % 4]
|
|
subnet["SubnetId"] = f"subnet-{idx:07d}"
|
|
return subnet
|
|
|
|
|
|
TWENTY_SUBNETS_IN_DIFFERENT_AZS = [subnet_in_az(i) for i in range(20)]
|
|
|
|
# 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._private.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_IN_BOUND_RULES = [
|
|
{
|
|
"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 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).
|
|
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"]
|
|
|
|
# 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
|
|
)
|
|
|
|
# Default launch template to expose to tests.
|
|
DEFAULT_LT = {
|
|
"LaunchTemplateId": "lt-00000000000000000",
|
|
"LaunchTemplateName": "ExampleLaunchTemplate",
|
|
"VersionNumber": 2,
|
|
"CreateTime": datetime(2020, 8, 17, 23, 30, 3),
|
|
"CreatedBy": DEFAULT_INSTANCE_PROFILE["Roles"][0]["Arn"],
|
|
"DefaultVersion": True,
|
|
"LaunchTemplateData": {
|
|
"EbsOptimized": False,
|
|
"IamInstanceProfile": {"Arn": DEFAULT_INSTANCE_PROFILE["Arn"]},
|
|
"NetworkInterfaces": [
|
|
{
|
|
"DeviceIndex": 0,
|
|
"Groups": [DEFAULT_SG["GroupId"]],
|
|
"SubnetId": DEFAULT_SUBNET["SubnetId"],
|
|
}
|
|
],
|
|
"ImageId": "ami-00000000000000000",
|
|
"InstanceType": "m5.large",
|
|
"TagSpecifications": [
|
|
{
|
|
"ResourceType": "instance",
|
|
"Tags": [{"Key": "test-key", "Value": "test-value"}],
|
|
},
|
|
{
|
|
"ResourceType": "volume",
|
|
"Tags": [{"Key": "test-key", "Value": "test-value"}],
|
|
},
|
|
],
|
|
},
|
|
}
|
|
|
|
# Default node provider tags to expose to tests.
|
|
DEFAULT_NODE_PROVIDER_INSTANCE_TAGS = {
|
|
TAG_RAY_NODE_KIND: NODE_KIND_HEAD,
|
|
TAG_RAY_LAUNCH_CONFIG: "test-ray-launch-config",
|
|
TAG_RAY_USER_NODE_TYPE: "ray.head.default",
|
|
}
|