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}}, )