def __init__(self, provider_config, node_ids, cluster_name): self.node_ids = set(node_ids) self.cluster_name = cluster_name self.provider_config = provider_config region = provider_config["region"] self.ec2_client = client_cache("ec2", region) self.ssm_client = client_cache("ssm", region) self.cloudwatch_client = client_cache("cloudwatch", region)
def __init__(self, provider_config, node_ids, cluster_name): # dedupe and sort node IDs to support deterministic unit test stubs self.node_ids = sorted(list(set(node_ids))) self.cluster_name = cluster_name self.provider_config = provider_config region = provider_config["region"] self.ec2_client = client_cache("ec2", region) self.ssm_client = client_cache("ssm", region) self.cloudwatch_client = client_cache("cloudwatch", region)
def __init__(self, provider_config: Dict[str, Any], node_id: str, cluster_name: str) -> None: self.node_id = node_id self.cluster_name = cluster_name self.provider_config = provider_config region = provider_config["region"] self.ec2_resource = resource_cache("ec2", region) self.ec2_client = self.ec2_resource.meta.client self.ssm_client = client_cache("ssm", region) cloudwatch_resource = resource_cache("cloudwatch", region) self.cloudwatch_client = cloudwatch_resource.meta.client self.CLOUDWATCH_CONFIG_TYPE_TO_CONFIG_VARIABLE_REPLACE_FUNC = { CloudwatchConfigType.AGENT: self._replace_cwa_config_variables, CloudwatchConfigType.DASHBOARD: self._replace_dashboard_config_variables, CloudwatchConfigType.ALARM: self._load_config_file, } self.CLOUDWATCH_CONFIG_TYPE_TO_UPDATE_FUNC_HEAD_NODE = { CloudwatchConfigType.AGENT: self._restart_cloudwatch_agent, CloudwatchConfigType.DASHBOARD: self._put_cloudwatch_dashboard, CloudwatchConfigType.ALARM: self._put_cloudwatch_alarm, } self.CLOUDWATCH_CONFIG_TYPE_TO_UPDATE_FUNC_WORKER_NODE = { CloudwatchConfigType.AGENT: self._restart_cloudwatch_agent, CloudwatchConfigType.ALARM: self._put_cloudwatch_alarm, }
def list_ec2_instances( region: str, aws_credentials: Dict[str, Any] = None) -> List[Dict[str, Any]]: """Get all instance-types/resources available in the user's AWS region. Args: region (str): the region of the AWS provider. e.g., "us-west-2". Returns: final_instance_types: a list of instances. An example of one element in the list: {'InstanceType': 'm5a.xlarge', 'ProcessorInfo': {'SupportedArchitectures': ['x86_64'], 'SustainedClockSpeedInGhz': 2.5},'VCpuInfo': {'DefaultVCpus': 4, 'DefaultCores': 2, 'DefaultThreadsPerCore': 2, 'ValidCores': [2], 'ValidThreadsPerCore': [1, 2]}, 'MemoryInfo': {'SizeInMiB': 16384}, ...} """ final_instance_types = [] aws_credentials = aws_credentials or {} ec2 = client_cache("ec2", region, BOTO_MAX_RETRIES, **aws_credentials) instance_types = ec2.describe_instance_types() final_instance_types.extend(copy.deepcopy(instance_types["InstanceTypes"])) while "NextToken" in instance_types: instance_types = ec2.describe_instance_types( NextToken=instance_types["NextToken"]) final_instance_types.extend( copy.deepcopy(instance_types["InstanceTypes"])) return final_instance_types
def resolve_policy_arns(config: Dict[str, Any], iam: Any, default_policy_arns: List[str]) -> List[str]: """Attach necessary AWS policies for CloudWatch related operations. Args: config: provider section of cluster config file. iam: AWS iam resource. default_policy_arns: List of default ray AWS policies. Returns: list of policy arns including additional policies for CloudWatch related operations if cloudwatch agent config is specifed in cluster config file. """ cwa_cfg_exists = CloudwatchHelper.cloudwatch_config_exists( config, "agent") if cwa_cfg_exists: cloudwatch_managed_policy = { "Version": "2012-10-17", "Statement": [{ "Effect": "Allow", "Action": [ "ssm:SendCommand", "ssm:ListCommandInvocations", "iam:PassRole", ], "Resource": "*", }], } iam_client = iam.meta.client iam_client.create_policy( PolicyName="CloudwatchManagedPolicies", PolicyDocument=json.dumps(cloudwatch_managed_policy), ) sts_client = client_cache("sts", config["region"]) account_id = sts_client.get_caller_identity().get("Account") managed_policy_arn = ( "arn:aws:iam::{}:policy/CloudwatchManagedPolicies".format( account_id)) policy_waiter = iam_client.get_waiter("policy_exists") policy_waiter.wait( PolicyArn=managed_policy_arn, WaiterConfig={ "Delay": 2, "MaxAttempts": 200 }, ) new_policy_arns = copy.copy(default_policy_arns) new_policy_arns.extend([ "arn:aws:iam::aws:policy/CloudWatchAgentAdminPolicy", "arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore", managed_policy_arn, ]) return new_policy_arns else: return default_policy_arns
def __init__(self, provider_config: Dict[str, Any], node_id: str, cluster_name: str) -> None: self.node_id = node_id self.cluster_name = cluster_name self.provider_config = provider_config region = provider_config["region"] self.ec2_resource = resource_cache("ec2", region) self.ec2_client = self.ec2_resource.meta.client self.ssm_client = client_cache("ssm", region) cloudwatch_resource = resource_cache("cloudwatch", region) self.cloudwatch_client = cloudwatch_resource.meta.client
def __init__(self, provider_config: Dict[str, Any], node_ids: List[str], cluster_name: str) -> None: # dedupe and sort node IDs to support deterministic unit test stubs self.node_ids = sorted(set(node_ids)) self.cluster_name = cluster_name self.provider_config = provider_config region = provider_config["region"] self.ec2_resource = resource_cache("ec2", region) self.ec2_client = self.ec2_resource.meta.client self.ssm_client = client_cache("ssm", region) cloudwatch_resource = resource_cache("cloudwatch", region) self.cloudwatch_client = cloudwatch_resource.meta.client
def ssm_client_stub(): client = client_cache("ssm", "us-west-2") with Stubber(client) as stubber: yield stubber stubber.assert_no_pending_responses()