wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/tests/aws/utils/constants.py
+228
View File
@@ -0,0 +1,228 @@
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",
}
python/ray/tests/aws/utils/helpers.py
+150
View File
@@ -0,0 +1,150 @@
import copy
import os
from typing import Any, Dict
import yaml
import ray
from ray.autoscaler._private.aws.cloudwatch.cloudwatch_helper import CloudwatchHelper
from ray.autoscaler._private.aws.node_provider import AWSNodeProvider
from ray.autoscaler._private.commands import prepare_config, validate_config
from ray.autoscaler.tags import (
NODE_KIND_HEAD,
NODE_KIND_WORKER,
TAG_RAY_CLUSTER_NAME,
TAG_RAY_NODE_KIND,
TAG_RAY_USER_NODE_TYPE,
)
from ray.tests.aws.utils.constants import (
DEFAULT_CLUSTER_NAME,
DEFAULT_NODE_PROVIDER_INSTANCE_TAGS,
)
def get_aws_example_config_file_path(file_name):
import ray.autoscaler.aws
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 fake_fillout_available_node_types_resources(config: Dict[str, Any]) -> None:
"""A cheap way to fill out the resources field (the same way a node
provider would autodetect them) as far as schema validation is concerned."""
available_node_types = config.get("available_node_types", {})
for label, value in available_node_types.items():
value["resources"] = value.get("resources", {"filler": 1})
def bootstrap_aws_config(config):
config = prepare_config(config)
fake_fillout_available_node_types_resources(config)
validate_config(config)
config["cluster_name"] = DEFAULT_CLUSTER_NAME
return ray.autoscaler._private.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)
def node_provider_tags(config: dict, type_name: str) -> dict:
"""
Returns a copy of DEFAULT_NODE_PROVIDER_INSTANCE_TAGS with the Ray node
kind and Ray user node type filled in from the input config and node type
name.
Args:
config: autoscaler config
type_name: node type name
Returns:
tags: node provider tags
"""
tags = copy.copy(DEFAULT_NODE_PROVIDER_INSTANCE_TAGS)
head_name = config["head_node_type"]
node_kind = NODE_KIND_HEAD if type_name is head_name else NODE_KIND_WORKER
tags[TAG_RAY_NODE_KIND] = node_kind
tags[TAG_RAY_USER_NODE_TYPE] = type_name
return tags
def apply_node_provider_config_updates(
config: dict, node_cfg: dict, node_type_name: str, max_count: int
) -> None:
"""
Applies default updates made by AWSNodeProvider to node_cfg during node
creation. This should only be used for testing purposes.
Args:
config: autoscaler config
node_cfg: node config
node_type_name: node type name
max_count: max nodes of the given type to launch
"""
tags = node_provider_tags(config, node_type_name)
tags[TAG_RAY_CLUSTER_NAME] = DEFAULT_CLUSTER_NAME
user_tag_specs = node_cfg.get("TagSpecifications", [])
tag_specs = [
{
"ResourceType": "instance",
"Tags": [{"Key": k, "Value": v} for k, v in sorted(tags.items())],
}
]
node_provider_cfg_updates = {
"MinCount": 1,
"MaxCount": max_count,
"TagSpecifications": tag_specs,
}
tags.pop(TAG_RAY_CLUSTER_NAME)
node_cfg.update(node_provider_cfg_updates)
# merge node provider tag specs with user overrides
AWSNodeProvider._merge_tag_specs(tag_specs, user_tag_specs)
def get_cloudwatch_agent_config_file_path():
return get_aws_example_config_file_path(
"cloudwatch/example-cloudwatch-agent-config.json"
)
def get_cloudwatch_dashboard_config_file_path():
return get_aws_example_config_file_path(
"cloudwatch/example-cloudwatch-dashboard-config.json"
)
def get_cloudwatch_alarm_config_file_path():
return get_aws_example_config_file_path(
"cloudwatch/example-cloudwatch-alarm-config.json"
)
def load_cloudwatch_example_config_file():
config = load_aws_example_config_file("example-cloudwatch.yaml")
cw_cfg = config["provider"]["cloudwatch"]
cw_cfg["agent"]["config"] = get_cloudwatch_agent_config_file_path()
cw_cfg["dashboard"]["config"] = get_cloudwatch_dashboard_config_file_path()
cw_cfg["alarm"]["config"] = get_cloudwatch_alarm_config_file_path()
return config
def get_cloudwatch_helper(node_ids):
config = load_cloudwatch_example_config_file()
config["cluster_name"] = DEFAULT_CLUSTER_NAME
return CloudwatchHelper(
config["provider"],
node_ids,
config["cluster_name"],
)
def get_ssm_param_name(cluster_name, config_type):
ssm_config_param_name = "AmazonCloudWatch-" + "ray_{}_config_{}".format(
config_type, cluster_name
)
return ssm_config_param_name
python/ray/tests/aws/utils/mocks.py
+7
View File
@@ -0,0 +1,7 @@
from ray.autoscaler._private.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
python/ray/tests/aws/utils/stubs.py
+584
View File
@@ -0,0 +1,584 @@
import copy
import json
from typing import Dict, List
from unittest import mock
from uuid import uuid4
from botocore.stub import ANY
import ray
from ray.autoscaler._private.aws.cloudwatch.cloudwatch_helper import (
CLOUDWATCH_AGENT_INSTALLED_TAG,
CLOUDWATCH_CONFIG_HASH_TAG_BASE,
)
from ray.autoscaler._private.aws.config import key_pair
from ray.autoscaler.tags import NODE_KIND_HEAD, TAG_RAY_NODE_KIND
from ray.tests.aws.utils import helpers
from ray.tests.aws.utils.constants import (
A_THOUSAND_SUBNETS_IN_DIFFERENT_VPCS,
DEFAULT_CLUSTER_NAME,
DEFAULT_INSTANCE_PROFILE,
DEFAULT_KEY_PAIR,
DEFAULT_LT,
DEFAULT_SUBNET,
TWENTY_SUBNETS_IN_DIFFERENT_AZS,
)
from ray.tests.aws.utils.helpers import (
get_cloudwatch_alarm_config_file_path,
get_cloudwatch_dashboard_config_file_path,
)
def configure_iam_role_default(iam_client_stub):
iam_client_stub.add_response(
"get_instance_profile",
expected_params={
"InstanceProfileName": ray.autoscaler._private.aws.config.DEFAULT_RAY_INSTANCE_PROFILE # noqa: E501
},
service_response={"InstanceProfile": DEFAULT_INSTANCE_PROFILE},
)
def configure_key_pair_default(
ec2_client_stub, region="us-west-2", expected_key_pair=DEFAULT_KEY_PAIR
):
patcher = mock.patch("os.path.exists")
def mock_path_exists_key_pair(path):
_, key_path = key_pair(0, region, expected_key_pair["KeyName"])
return path == key_path
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": [expected_key_pair["KeyName"]]}]
},
service_response={"KeyPairs": [expected_key_pair]},
)
def configure_subnet_default(ec2_client_stub):
ec2_client_stub.add_response(
"describe_subnets",
expected_params={},
service_response={"Subnets": [DEFAULT_SUBNET]},
)
def describe_a_thousand_subnets_in_different_vpcs(ec2_client_stub):
ec2_client_stub.add_response(
"describe_subnets",
expected_params={},
service_response={"Subnets": A_THOUSAND_SUBNETS_IN_DIFFERENT_VPCS},
)
def describe_twenty_subnets_in_different_azs(ec2_client_stub):
ec2_client_stub.add_response(
"describe_subnets",
expected_params={},
service_response={"Subnets": TWENTY_SUBNETS_IN_DIFFERENT_AZS},
)
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, subnets: List[Dict[str, str]]):
ec2_client_stub.add_response(
"describe_subnets",
expected_params={
"Filters": [
{"Name": "subnet-id", "Values": [s["SubnetId"] for s in subnets]}
]
},
service_response={"Subnets": subnets},
)
def describe_no_security_groups(ec2_client_stub):
ec2_client_stub.add_response(
"describe_security_groups",
expected_params={"Filters": ANY},
service_response={},
)
def describe_a_security_group(ec2_client_stub, security_group):
ec2_client_stub.add_response(
"describe_security_groups",
expected_params={
"Filters": [{"Name": "group-id", "Values": [security_group["GroupId"]]}]
},
service_response={"SecurityGroups": [security_group]},
)
def describe_an_sg_2(ec2_client_stub, security_group):
"""Same as last function, different input param format.
A call with this input parameter format is made when sg.ip_permissions is
accessed in aws/config.py.
"""
ec2_client_stub.add_response(
"describe_security_groups",
expected_params={"GroupIds": [security_group["GroupId"]]},
service_response={"SecurityGroups": [security_group]},
)
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"],
"TagSpecifications": [
{
"ResourceType": "security-group",
"Tags": [
{
"Key": ray.autoscaler._private.aws.config.RAY,
"Value": "true",
},
{"Key": "ray-cluster-name", "Value": DEFAULT_CLUSTER_NAME},
],
},
],
},
service_response={"GroupId": security_group["GroupId"]},
)
def describe_sgs_by_id(ec2_client_stub, security_group_ids, security_groups):
ec2_client_stub.add_response(
"describe_security_groups",
expected_params={
"Filters": [{"Name": "group-id", "Values": security_group_ids}]
},
service_response={"SecurityGroups": security_groups},
)
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]},
)
def run_instances_with_network_interfaces_consumer(ec2_client_stub, network_interfaces):
ec2_client_stub.add_response(
"run_instances",
expected_params={
"NetworkInterfaces": network_interfaces,
"ImageId": ANY,
"InstanceType": ANY,
"KeyName": ANY,
"MinCount": ANY,
"MaxCount": ANY,
"TagSpecifications": ANY,
},
service_response={},
)
def run_instances_with_launch_template_consumer(
ec2_client_stub, config, node_cfg, node_type_name, lt_data, max_count
):
# create a copy of both node config and launch template data to modify
lt_data_cp = copy.deepcopy(lt_data)
node_cfg_cp = copy.deepcopy(node_cfg)
# override launch template parameters with explicit node config parameters
lt_data_cp.update(node_cfg_cp)
# copy all launch template parameters back to node config
node_cfg_cp.update(lt_data_cp)
# copy all default node provider config updates to node config
helpers.apply_node_provider_config_updates(
config, node_cfg_cp, node_type_name, max_count
)
# remove any security group and subnet IDs copied from network interfaces
node_cfg_cp.pop("SecurityGroupIds", [])
node_cfg_cp.pop("SubnetIds", [])
ec2_client_stub.add_response(
"run_instances", expected_params=node_cfg_cp, service_response={}
)
def describe_instances_with_any_filter_consumer(ec2_client_stub):
ec2_client_stub.add_response(
"describe_instances", expected_params={"Filters": ANY}, service_response={}
)
def describe_launch_template_versions_by_id_default(ec2_client_stub, versions):
ec2_client_stub.add_response(
"describe_launch_template_versions",
expected_params={
"LaunchTemplateId": DEFAULT_LT["LaunchTemplateId"],
"Versions": versions,
},
service_response={"LaunchTemplateVersions": [DEFAULT_LT]},
)
def describe_launch_template_versions_by_name_default(ec2_client_stub, versions):
ec2_client_stub.add_response(
"describe_launch_template_versions",
expected_params={
"LaunchTemplateName": DEFAULT_LT["LaunchTemplateName"],
"Versions": versions,
},
service_response={"LaunchTemplateVersions": [DEFAULT_LT]},
)
def get_ec2_cwa_installed_tag_true(ec2_client_stub, node_id):
ec2_client_stub.add_response(
"describe_instances",
expected_params={"InstanceIds": [node_id]},
service_response={
"Reservations": [
{
"Instances": [
{
"InstanceId": node_id,
"Tags": [
{
"Key": CLOUDWATCH_AGENT_INSTALLED_TAG,
"Value": "True",
},
],
}
]
}
]
},
)
def update_hash_tag_success(ec2_client_stub, node_id, config_type, cloudwatch_helper):
hash_key_value = "-".join([CLOUDWATCH_CONFIG_HASH_TAG_BASE, config_type])
cur_hash_value = get_sha256_hash_of_cloudwatch_config_file(
config_type, cloudwatch_helper
)
ec2_client_stub.add_response(
"create_tags",
expected_params={
"Resources": [node_id],
"Tags": [{"Key": hash_key_value, "Value": cur_hash_value}],
},
service_response={"ResponseMetadata": {"HTTPStatusCode": 200}},
)
def add_cwa_installed_tag_response(ec2_client_stub, node_id):
ec2_client_stub.add_response(
"create_tags",
expected_params={
"Resources": node_id,
"Tags": [{"Key": CLOUDWATCH_AGENT_INSTALLED_TAG, "Value": "True"}],
},
service_response={"ResponseMetadata": {"HTTPStatusCode": 200}},
)
def get_head_node_config_hash_different(ec2_client_stub, config_type, cwh, node_id):
hash_key_value = "-".join([CLOUDWATCH_CONFIG_HASH_TAG_BASE, config_type])
cur_hash_value = get_sha256_hash_of_cloudwatch_config_file(config_type, cwh)
filters = cwh._get_current_cluster_session_nodes(cwh.cluster_name)
filters.append(
{
"Name": "tag:{}".format(TAG_RAY_NODE_KIND),
"Values": [NODE_KIND_HEAD],
}
)
ec2_client_stub.add_response(
"describe_instances",
expected_params={"Filters": filters},
service_response={
"Reservations": [
{
"Instances": [
{
"InstanceId": node_id,
"Tags": [
{"Key": hash_key_value, "Value": cur_hash_value},
],
}
]
}
]
},
)
def get_cur_node_config_hash_different(ec2_client_stub, config_type, node_id):
hash_key_value = "-".join([CLOUDWATCH_CONFIG_HASH_TAG_BASE, config_type])
ec2_client_stub.add_response(
"describe_instances",
expected_params={"InstanceIds": [node_id]},
service_response={
"Reservations": [
{
"Instances": [
{
"InstanceId": node_id,
"Tags": [
{"Key": hash_key_value, "Value": str(uuid4())},
],
}
]
}
]
},
)
def send_command_cwa_install(ssm_client_stub, node_id):
command_id = str(uuid4())
ssm_client_stub.add_response(
"send_command",
expected_params={
"DocumentName": "AWS-ConfigureAWSPackage",
"InstanceIds": node_id,
"MaxConcurrency": "1",
"MaxErrors": "0",
"Parameters": {
"action": ["Install"],
"name": ["AmazonCloudWatchAgent"],
"version": ["latest"],
},
},
service_response={
"Command": {
"CommandId": command_id,
"DocumentName": "AWS-ConfigureAWSPackage",
}
},
)
return command_id
def list_command_invocations_status(ssm_client_stub, node_id, cmd_id, status):
ssm_client_stub.add_response(
"list_command_invocations",
expected_params={"CommandId": cmd_id, "InstanceId": node_id},
service_response={"CommandInvocations": [{"Status": status}]},
)
def list_command_invocations_failed(ssm_client_stub, node_id, cmd_id):
status = "Failed"
list_command_invocations_status(ssm_client_stub, node_id, cmd_id, status)
def list_command_invocations_success(ssm_client_stub, node_id, cmd_id):
status = "Success"
list_command_invocations_status(ssm_client_stub, node_id, cmd_id, status)
def put_parameter_cloudwatch_config(ssm_client_stub, cluster_name, section_name):
ssm_config_param_name = helpers.get_ssm_param_name(cluster_name, section_name)
ssm_client_stub.add_response(
"put_parameter",
expected_params={
"Name": ssm_config_param_name,
"Type": "String",
"Value": ANY,
"Overwrite": True,
"Tier": ANY,
},
service_response={},
)
def send_command_cwa_collectd_init(ssm_client_stub, node_id):
command_id = str(uuid4())
ssm_client_stub.add_response(
"send_command",
expected_params={
"DocumentName": "AWS-RunShellScript",
"InstanceIds": [node_id],
"MaxConcurrency": "1",
"MaxErrors": "0",
"Parameters": {
"commands": [
"mkdir -p /usr/share/collectd/",
"touch /usr/share/collectd/types.db",
],
},
},
service_response={"Command": {"CommandId": command_id}},
)
return command_id
def send_command_start_cwa(ssm_client_stub, node_id, parameter_name):
command_id = str(uuid4())
ssm_client_stub.add_response(
"send_command",
expected_params={
"DocumentName": "AmazonCloudWatch-ManageAgent",
"InstanceIds": [node_id],
"MaxConcurrency": "1",
"MaxErrors": "0",
"Parameters": {
"action": ["configure"],
"mode": ["ec2"],
"optionalConfigurationSource": ["ssm"],
"optionalConfigurationLocation": [parameter_name],
"optionalRestart": ["yes"],
},
},
service_response={"Command": {"CommandId": command_id}},
)
return command_id
def send_command_stop_cwa(ssm_client_stub, node_id):
command_id = str(uuid4())
ssm_client_stub.add_response(
"send_command",
expected_params={
"DocumentName": "AmazonCloudWatch-ManageAgent",
"InstanceIds": [node_id],
"MaxConcurrency": "1",
"MaxErrors": "0",
"Parameters": {
"action": ["stop"],
"mode": ["ec2"],
},
},
service_response={"Command": {"CommandId": command_id}},
)
return command_id
def get_param_ssm_same(ssm_client_stub, ssm_param_name, cloudwatch_helper, config_type):
command_id = str(uuid4())
cw_value_json = (
cloudwatch_helper.CLOUDWATCH_CONFIG_TYPE_TO_CONFIG_VARIABLE_REPLACE_FUNC.get(
config_type
)(config_type)
)
ssm_client_stub.add_response(
"get_parameter",
expected_params={"Name": ssm_param_name},
service_response={"Parameter": {"Value": json.dumps(cw_value_json)}},
)
return command_id
def get_sha256_hash_of_cloudwatch_config_file(config_type, cloudwatch_helper):
cw_value_file = cloudwatch_helper._sha256_hash_file(config_type)
return cw_value_file
def get_param_ssm_different(ssm_client_stub, ssm_param_name):
command_id = str(uuid4())
ssm_client_stub.add_response(
"get_parameter",
expected_params={"Name": ssm_param_name},
service_response={"Parameter": {"Value": "value"}},
)
return command_id
def get_param_ssm_exception(ssm_client_stub, ssm_param_name):
command_id = str(uuid4())
ssm_client_stub.add_client_error(
"get_parameter",
"ParameterNotFound",
expected_params={"Name": ssm_param_name},
response_meta={"Error": {"Code": "ParameterNotFound"}},
)
return command_id
def put_cluster_dashboard_success(cloudwatch_client_stub, cloudwatch_helper):
widgets = []
json_config_path = get_cloudwatch_dashboard_config_file_path()
with open(json_config_path) as f:
dashboard_config = json.load(f)
for item in dashboard_config:
item_out = cloudwatch_helper._replace_all_config_variables(
item,
cloudwatch_helper.node_id,
cloudwatch_helper.cluster_name,
cloudwatch_helper.provider_config["region"],
)
widgets.append(item_out)
dashboard_name = cloudwatch_helper.cluster_name + "-" + "example-dashboard-name"
cloudwatch_client_stub.add_response(
"put_dashboard",
expected_params={
"DashboardName": dashboard_name,
"DashboardBody": json.dumps({"widgets": widgets}),
},
service_response={"ResponseMetadata": {"HTTPStatusCode": 200}},
)
def put_cluster_alarms_success(cloudwatch_client_stub, cloudwatch_helper):
json_config_path = get_cloudwatch_alarm_config_file_path()
with open(json_config_path) as f:
data = json.load(f)
for item in data:
item_out = copy.deepcopy(item)
cloudwatch_helper._replace_all_config_variables(
item_out,
cloudwatch_helper.node_id,
cloudwatch_helper.cluster_name,
cloudwatch_helper.provider_config["region"],
)
cloudwatch_client_stub.add_response(
"put_metric_alarm",
expected_params=item_out,
service_response={"ResponseMetadata": {"HTTPStatusCode": 200}},
)
def get_metric_alarm(cloudwatch_client_stub):
cloudwatch_client_stub.add_response(
"describe_alarms",
expected_params={},
service_response={"MetricAlarms": [{"AlarmName": "myalarm"}]},
)
def delete_metric_alarms(cloudwatch_client_stub):
cloudwatch_client_stub.add_response(
"delete_alarms",
expected_params={"AlarmNames": ["myalarm"]},
service_response={"ResponseMetadata": {"HTTPStatusCode": 200}},
)