chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,821 @@
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from enum import Enum
|
||||
from typing import Any, Callable, Dict, List, Union
|
||||
|
||||
import botocore
|
||||
|
||||
from ray.autoscaler._private.aws.utils import client_cache, resource_cache
|
||||
from ray.autoscaler.tags import NODE_KIND_HEAD, TAG_RAY_CLUSTER_NAME, TAG_RAY_NODE_KIND
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
RAY = "ray-autoscaler"
|
||||
CLOUDWATCH_RAY_INSTANCE_PROFILE = RAY + "-cloudwatch-v1"
|
||||
CLOUDWATCH_RAY_IAM_ROLE = RAY + "-cloudwatch-v1"
|
||||
CLOUDWATCH_AGENT_INSTALLED_AMI_TAG = "T6Iq2faj"
|
||||
CLOUDWATCH_AGENT_INSTALLED_TAG = "cloudwatch-agent-installed"
|
||||
CLOUDWATCH_CONFIG_HASH_TAG_BASE = "cloudwatch-config-hash"
|
||||
|
||||
|
||||
class CloudwatchConfigType(str, Enum):
|
||||
AGENT = "agent"
|
||||
DASHBOARD = "dashboard"
|
||||
ALARM = "alarm"
|
||||
|
||||
|
||||
class CloudwatchHelper:
|
||||
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: Dict[
|
||||
str, Callable
|
||||
] = {
|
||||
CloudwatchConfigType.AGENT.value: self._replace_cwa_config_vars,
|
||||
CloudwatchConfigType.DASHBOARD.value: self._replace_dashboard_config_vars,
|
||||
CloudwatchConfigType.ALARM.value: self._load_config_file,
|
||||
}
|
||||
self.CLOUDWATCH_CONFIG_TYPE_TO_UPDATE_FUNC_HEAD_NODE: Dict[str, Callable] = {
|
||||
CloudwatchConfigType.AGENT.value: self._restart_cloudwatch_agent,
|
||||
CloudwatchConfigType.DASHBOARD.value: self._put_cloudwatch_dashboard,
|
||||
CloudwatchConfigType.ALARM.value: self._put_cloudwatch_alarm,
|
||||
}
|
||||
self.CLOUDWATCH_CONFIG_TYPE_TO_UPDATE_FUNC_WORKER_NODE: Dict[str, Callable] = {
|
||||
CloudwatchConfigType.AGENT.value: self._restart_cloudwatch_agent,
|
||||
CloudwatchConfigType.ALARM.value: self._put_cloudwatch_alarm,
|
||||
}
|
||||
|
||||
def update_from_config(self, is_head_node: bool) -> None:
|
||||
"""Discovers and applies CloudWatch config updates as required.
|
||||
|
||||
Args:
|
||||
is_head_node: whether this node is the head node.
|
||||
"""
|
||||
for config_type in CloudwatchConfigType:
|
||||
if CloudwatchHelper.cloudwatch_config_exists(
|
||||
self.provider_config, config_type.value
|
||||
):
|
||||
self._update_cloudwatch_config(config_type.value, is_head_node)
|
||||
|
||||
def _ec2_health_check_waiter(self, node_id: str) -> None:
|
||||
# wait for all EC2 instance checks to complete
|
||||
try:
|
||||
logger.info(
|
||||
"Waiting for EC2 instance health checks to complete before "
|
||||
"configuring Unified Cloudwatch Agent. This may take a few "
|
||||
"minutes..."
|
||||
)
|
||||
waiter = self.ec2_client.get_waiter("instance_status_ok")
|
||||
waiter.wait(InstanceIds=[node_id])
|
||||
except botocore.exceptions.WaiterError as e:
|
||||
logger.error(
|
||||
"Failed while waiting for EC2 instance checks to complete: {}".format(
|
||||
e.message
|
||||
)
|
||||
)
|
||||
raise e
|
||||
|
||||
def _update_cloudwatch_config(self, config_type: str, is_head_node: bool) -> None:
|
||||
"""
|
||||
check whether update operations are needed in
|
||||
cloudwatch related configs
|
||||
"""
|
||||
cwa_installed = self._setup_cwa()
|
||||
param_name = self._get_ssm_param_name(config_type)
|
||||
if cwa_installed:
|
||||
if is_head_node:
|
||||
cw_config_ssm = self._set_cloudwatch_ssm_config_param(
|
||||
param_name, config_type
|
||||
)
|
||||
cur_cw_config_hash = self._sha256_hash_file(config_type)
|
||||
ssm_cw_config_hash = self._sha256_hash_json(cw_config_ssm)
|
||||
# check if user updated cloudwatch related config files.
|
||||
# if so, perform corresponding actions.
|
||||
if cur_cw_config_hash != ssm_cw_config_hash:
|
||||
logger.info(
|
||||
"Cloudwatch {} config file has changed.".format(config_type)
|
||||
)
|
||||
self._upload_config_to_ssm_and_set_hash_tag(config_type)
|
||||
self.CLOUDWATCH_CONFIG_TYPE_TO_UPDATE_FUNC_HEAD_NODE.get(
|
||||
config_type
|
||||
)()
|
||||
else:
|
||||
head_node_hash = self._get_head_node_config_hash(config_type)
|
||||
cur_node_hash = self._get_cur_node_config_hash(config_type)
|
||||
if head_node_hash != cur_node_hash:
|
||||
logger.info(
|
||||
"Cloudwatch {} config file has changed.".format(config_type)
|
||||
)
|
||||
update_func = (
|
||||
self.CLOUDWATCH_CONFIG_TYPE_TO_UPDATE_FUNC_WORKER_NODE.get(
|
||||
config_type
|
||||
)
|
||||
)
|
||||
if update_func:
|
||||
update_func()
|
||||
self._update_cloudwatch_hash_tag_value(
|
||||
self.node_id, head_node_hash, config_type
|
||||
)
|
||||
|
||||
def _put_cloudwatch_dashboard(self) -> Dict[str, Any]:
|
||||
"""put dashboard to cloudwatch console"""
|
||||
|
||||
cloudwatch_config = self.provider_config["cloudwatch"]
|
||||
dashboard_config = cloudwatch_config.get("dashboard", {})
|
||||
dashboard_name_cluster = dashboard_config.get("name", self.cluster_name)
|
||||
dashboard_name = self.cluster_name + "-" + dashboard_name_cluster
|
||||
|
||||
widgets = self._replace_dashboard_config_vars(
|
||||
CloudwatchConfigType.DASHBOARD.value
|
||||
)
|
||||
|
||||
response = self.cloudwatch_client.put_dashboard(
|
||||
DashboardName=dashboard_name, DashboardBody=json.dumps({"widgets": widgets})
|
||||
)
|
||||
issue_count = len(response.get("DashboardValidationMessages", []))
|
||||
if issue_count > 0:
|
||||
for issue in response.get("DashboardValidationMessages"):
|
||||
logging.error(
|
||||
"Error in dashboard config: {} - {}".format(
|
||||
issue["Message"], issue["DataPath"]
|
||||
)
|
||||
)
|
||||
raise Exception(
|
||||
"Errors in dashboard configuration: {} issues raised".format(
|
||||
issue_count
|
||||
)
|
||||
)
|
||||
else:
|
||||
logger.info("Successfully put dashboard to CloudWatch console")
|
||||
return response
|
||||
|
||||
def _put_cloudwatch_alarm(self) -> None:
|
||||
"""put CloudWatch metric alarms read from config"""
|
||||
param_name = self._get_ssm_param_name(CloudwatchConfigType.ALARM.value)
|
||||
data = json.loads(self._get_ssm_param(param_name))
|
||||
for item in data:
|
||||
item_out = copy.deepcopy(item)
|
||||
self._replace_all_config_variables(
|
||||
item_out,
|
||||
self.node_id,
|
||||
self.cluster_name,
|
||||
self.provider_config["region"],
|
||||
)
|
||||
self.cloudwatch_client.put_metric_alarm(**item_out)
|
||||
logger.info("Successfully put alarms to CloudWatch console")
|
||||
|
||||
def _send_command_to_node(
|
||||
self, document_name: str, parameters: Dict[str, List[str]], node_id: str
|
||||
) -> Dict[str, Any]:
|
||||
"""send SSM command to the given nodes"""
|
||||
logger.debug(
|
||||
"Sending SSM command to {} node(s). Document name: {}. "
|
||||
"Parameters: {}.".format(node_id, document_name, parameters)
|
||||
)
|
||||
response = self.ssm_client.send_command(
|
||||
InstanceIds=[node_id],
|
||||
DocumentName=document_name,
|
||||
Parameters=parameters,
|
||||
MaxConcurrency="1",
|
||||
MaxErrors="0",
|
||||
)
|
||||
return response
|
||||
|
||||
def _ssm_command_waiter(
|
||||
self,
|
||||
document_name: str,
|
||||
parameters: Dict[str, List[str]],
|
||||
node_id: str,
|
||||
retry_failed: bool = True,
|
||||
) -> Dict[str, Any]:
|
||||
"""wait for SSM command to complete on all cluster nodes"""
|
||||
|
||||
# This waiter differs from the built-in SSM.Waiter by
|
||||
# optimistically waiting for the command invocation to
|
||||
# exist instead of failing immediately, and by resubmitting
|
||||
# any failed command until all retry attempts are exhausted
|
||||
# by default.
|
||||
response = self._send_command_to_node(document_name, parameters, node_id)
|
||||
command_id = response["Command"]["CommandId"]
|
||||
|
||||
cloudwatch_config = self.provider_config["cloudwatch"]
|
||||
agent_retryer_config = cloudwatch_config.get(
|
||||
CloudwatchConfigType.AGENT.value
|
||||
).get("retryer", {})
|
||||
max_attempts = agent_retryer_config.get("max_attempts", 120)
|
||||
delay_seconds = agent_retryer_config.get("delay_seconds", 30)
|
||||
num_attempts = 0
|
||||
cmd_invocation_res = {}
|
||||
while True:
|
||||
num_attempts += 1
|
||||
logger.debug(
|
||||
"Listing SSM command ID {} invocations on node {}".format(
|
||||
command_id, node_id
|
||||
)
|
||||
)
|
||||
response = self.ssm_client.list_command_invocations(
|
||||
CommandId=command_id,
|
||||
InstanceId=node_id,
|
||||
)
|
||||
cmd_invocations = response["CommandInvocations"]
|
||||
if not cmd_invocations:
|
||||
logger.debug(
|
||||
"SSM Command ID {} invocation does not exist. If "
|
||||
"the command was just started, it may take a "
|
||||
"few seconds to register.".format(command_id)
|
||||
)
|
||||
else:
|
||||
if len(cmd_invocations) > 1:
|
||||
logger.warning(
|
||||
"Expected to find 1 SSM command invocation with "
|
||||
"ID {} on node {} but found {}: {}".format(
|
||||
command_id,
|
||||
node_id,
|
||||
len(cmd_invocations),
|
||||
cmd_invocations,
|
||||
)
|
||||
)
|
||||
cmd_invocation = cmd_invocations[0]
|
||||
if cmd_invocation["Status"] == "Success":
|
||||
logger.debug(
|
||||
"SSM Command ID {} completed successfully.".format(command_id)
|
||||
)
|
||||
cmd_invocation_res[node_id] = True
|
||||
break
|
||||
if num_attempts >= max_attempts:
|
||||
logger.error(
|
||||
"Max attempts for command {} exceeded on node {}".format(
|
||||
command_id, node_id
|
||||
)
|
||||
)
|
||||
raise botocore.exceptions.WaiterError(
|
||||
name="ssm_waiter",
|
||||
reason="Max attempts exceeded",
|
||||
last_response=cmd_invocation,
|
||||
)
|
||||
if cmd_invocation["Status"] == "Failed":
|
||||
logger.debug(f"SSM Command ID {command_id} failed.")
|
||||
if retry_failed:
|
||||
logger.debug(f"Retrying in {delay_seconds} seconds.")
|
||||
response = self._send_command_to_node(
|
||||
document_name, parameters, node_id
|
||||
)
|
||||
command_id = response["Command"]["CommandId"]
|
||||
logger.debug(
|
||||
"Sent SSM command ID {} to node {}".format(
|
||||
command_id, node_id
|
||||
)
|
||||
)
|
||||
else:
|
||||
logger.debug(f"Ignoring Command ID {command_id} failure.")
|
||||
cmd_invocation_res[node_id] = False
|
||||
break
|
||||
time.sleep(delay_seconds)
|
||||
|
||||
return cmd_invocation_res
|
||||
|
||||
def _replace_config_variables(
|
||||
self, string: str, node_id: str, cluster_name: str, region: str
|
||||
) -> str:
|
||||
"""
|
||||
replace known config variable occurrences in the input string
|
||||
does not replace variables with undefined or empty strings
|
||||
"""
|
||||
|
||||
if node_id:
|
||||
string = string.replace("{instance_id}", node_id)
|
||||
if cluster_name:
|
||||
string = string.replace("{cluster_name}", cluster_name)
|
||||
if region:
|
||||
string = string.replace("{region}", region)
|
||||
return string
|
||||
|
||||
def _replace_all_config_variables(
|
||||
self,
|
||||
collection: Union[Dict[str, Any], str],
|
||||
node_id: str,
|
||||
cluster_name: str,
|
||||
region: str,
|
||||
) -> Union[str, Dict[str, Any]]:
|
||||
"""
|
||||
Replace known config variable occurrences in the input collection.
|
||||
The input collection must be either a dict or list.
|
||||
Returns a tuple consisting of the output collection and the number of
|
||||
modified strings in the collection (which is not necessarily equal to
|
||||
the number of variables replaced).
|
||||
"""
|
||||
|
||||
for key in collection:
|
||||
if type(collection) is dict:
|
||||
value = collection.get(key)
|
||||
index_key = key
|
||||
elif type(collection) is list:
|
||||
value = key
|
||||
index_key = collection.index(key)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Can't replace CloudWatch config variables "
|
||||
f"in unsupported collection type: {type(collection)}."
|
||||
f"Please check your CloudWatch JSON config files."
|
||||
)
|
||||
if type(value) is str:
|
||||
collection[index_key] = self._replace_config_variables(
|
||||
value, node_id, cluster_name, region
|
||||
)
|
||||
elif type(value) is dict or type(value) is list:
|
||||
collection[index_key] = self._replace_all_config_variables(
|
||||
value, node_id, cluster_name, region
|
||||
)
|
||||
return collection
|
||||
|
||||
def _load_config_file(self, config_type: str) -> Dict[str, Any]:
|
||||
"""load JSON config file"""
|
||||
cloudwatch_config = self.provider_config["cloudwatch"]
|
||||
json_config_file_section = cloudwatch_config.get(config_type, {})
|
||||
json_config_file_path = json_config_file_section.get("config", {})
|
||||
json_config_path = os.path.abspath(json_config_file_path)
|
||||
with open(json_config_path) as f:
|
||||
data = json.load(f)
|
||||
return data
|
||||
|
||||
def _set_cloudwatch_ssm_config_param(
|
||||
self, parameter_name: str, config_type: str
|
||||
) -> str:
|
||||
"""
|
||||
get cloudwatch config for the given param and config type from SSM
|
||||
if it exists, put it in the SSM param store if not
|
||||
"""
|
||||
try:
|
||||
parameter_value = self._get_ssm_param(parameter_name)
|
||||
except botocore.exceptions.ClientError as e:
|
||||
if e.response["Error"]["Code"] == "ParameterNotFound":
|
||||
logger.info(
|
||||
"Cloudwatch {} config file is not found "
|
||||
"at SSM parameter store. "
|
||||
"Checking for Unified CloudWatch Agent installation".format(
|
||||
config_type
|
||||
)
|
||||
)
|
||||
return self._get_default_empty_config_file_hash()
|
||||
else:
|
||||
logger.info(
|
||||
"Failed to fetch Unified CloudWatch Agent config from SSM "
|
||||
"parameter store."
|
||||
)
|
||||
logger.error(e)
|
||||
raise e
|
||||
return parameter_value
|
||||
|
||||
def _get_default_empty_config_file_hash(self):
|
||||
default_cw_config = "{}"
|
||||
parameter_value = self._sha256_hash_json(default_cw_config)
|
||||
return parameter_value
|
||||
|
||||
def _get_ssm_param(self, parameter_name: str) -> str:
|
||||
"""
|
||||
get the SSM parameter value associated with the given parameter name
|
||||
"""
|
||||
response = self.ssm_client.get_parameter(Name=parameter_name)
|
||||
logger.info("Successfully fetch ssm parameter: {}".format(parameter_name))
|
||||
res = response.get("Parameter", {})
|
||||
cwa_parameter = res.get("Value", {})
|
||||
return cwa_parameter
|
||||
|
||||
def _sha256_hash_json(self, value: str) -> str:
|
||||
"""calculate the json string sha256 hash"""
|
||||
sha256_hash = hashlib.new("sha256")
|
||||
binary_value = value.encode("utf-8")
|
||||
sha256_hash.update(binary_value)
|
||||
sha256_res = sha256_hash.hexdigest()
|
||||
return sha256_res
|
||||
|
||||
def _sha256_hash_file(self, config_type: str) -> str:
|
||||
"""calculate the config file sha256 hash"""
|
||||
config = self.CLOUDWATCH_CONFIG_TYPE_TO_CONFIG_VARIABLE_REPLACE_FUNC.get(
|
||||
config_type
|
||||
)(config_type)
|
||||
value = json.dumps(config)
|
||||
sha256_res = self._sha256_hash_json(value)
|
||||
return sha256_res
|
||||
|
||||
def _upload_config_to_ssm_and_set_hash_tag(self, config_type: str):
|
||||
data = self.CLOUDWATCH_CONFIG_TYPE_TO_CONFIG_VARIABLE_REPLACE_FUNC.get(
|
||||
config_type
|
||||
)(config_type)
|
||||
sha256_hash_value = self._sha256_hash_file(config_type)
|
||||
self._upload_config_to_ssm(data, config_type)
|
||||
self._update_cloudwatch_hash_tag_value(
|
||||
self.node_id, sha256_hash_value, config_type
|
||||
)
|
||||
|
||||
def _add_cwa_installed_tag(self, node_id: str) -> None:
|
||||
self.ec2_client.create_tags(
|
||||
Resources=[node_id],
|
||||
Tags=[{"Key": CLOUDWATCH_AGENT_INSTALLED_TAG, "Value": "True"}],
|
||||
)
|
||||
logger.info(
|
||||
"Successfully add Unified CloudWatch Agent installed "
|
||||
"tag on {}".format(node_id)
|
||||
)
|
||||
|
||||
def _update_cloudwatch_hash_tag_value(
|
||||
self, node_id: str, sha256_hash_value: str, config_type: str
|
||||
):
|
||||
hash_key_value = "-".join([CLOUDWATCH_CONFIG_HASH_TAG_BASE, config_type])
|
||||
self.ec2_client.create_tags(
|
||||
Resources=[node_id],
|
||||
Tags=[{"Key": hash_key_value, "Value": sha256_hash_value}],
|
||||
)
|
||||
logger.info(
|
||||
"Successfully update cloudwatch {} hash tag on {}".format(
|
||||
config_type, node_id
|
||||
)
|
||||
)
|
||||
|
||||
def _get_ssm_param_name(self, config_type: str) -> str:
|
||||
"""return the parameter name for cloudwatch configs"""
|
||||
ssm_config_param_name = "AmazonCloudWatch-" + "ray_{}_config_{}".format(
|
||||
config_type, self.cluster_name
|
||||
)
|
||||
return ssm_config_param_name
|
||||
|
||||
def _put_ssm_param(self, parameter: Dict[str, Any], parameter_name: str) -> None:
|
||||
"""upload cloudwatch config to the SSM parameter store"""
|
||||
self.ssm_client.put_parameter(
|
||||
Name=parameter_name,
|
||||
Type="String",
|
||||
Value=json.dumps(parameter),
|
||||
Overwrite=True,
|
||||
Tier="Intelligent-Tiering",
|
||||
)
|
||||
|
||||
def _upload_config_to_ssm(self, param: Dict[str, Any], config_type: str):
|
||||
param_name = self._get_ssm_param_name(config_type)
|
||||
self._put_ssm_param(param, param_name)
|
||||
|
||||
def _replace_cwa_config_vars(self, config_type: str) -> Dict[str, Any]:
|
||||
"""
|
||||
replace {instance_id}, {region}, {cluster_name}
|
||||
variable occurrences in Unified Cloudwatch Agent config file
|
||||
"""
|
||||
cwa_config = self._load_config_file(config_type)
|
||||
self._replace_all_config_variables(
|
||||
cwa_config,
|
||||
self.node_id,
|
||||
self.cluster_name,
|
||||
self.provider_config["region"],
|
||||
)
|
||||
return cwa_config
|
||||
|
||||
def _replace_dashboard_config_vars(self, config_type: str) -> List[str]:
|
||||
"""
|
||||
replace known variable occurrences in CloudWatch Dashboard config file
|
||||
"""
|
||||
data = self._load_config_file(config_type)
|
||||
widgets = []
|
||||
for item in data:
|
||||
item_out = self._replace_all_config_variables(
|
||||
item,
|
||||
self.node_id,
|
||||
self.cluster_name,
|
||||
self.provider_config["region"],
|
||||
)
|
||||
widgets.append(item_out)
|
||||
return widgets
|
||||
|
||||
def _replace_alarm_config_vars(self, config_type: str) -> List[str]:
|
||||
"""
|
||||
replace {instance_id}, {region}, {cluster_name}
|
||||
variable occurrences in cloudwatch alarm config file
|
||||
"""
|
||||
data = self._load_config_file(config_type)
|
||||
param_data = []
|
||||
for item in data:
|
||||
item_out = copy.deepcopy(item)
|
||||
self._replace_all_config_variables(
|
||||
item_out,
|
||||
self.node_id,
|
||||
self.cluster_name,
|
||||
self.provider_config["region"],
|
||||
)
|
||||
param_data.append(item_out)
|
||||
return param_data
|
||||
|
||||
def _restart_cloudwatch_agent(self) -> None:
|
||||
"""restart Unified CloudWatch Agent"""
|
||||
cwa_param_name = self._get_ssm_param_name(CloudwatchConfigType.AGENT.value)
|
||||
logger.info(
|
||||
"Restarting Unified CloudWatch Agent package on node {}.".format(
|
||||
self.node_id
|
||||
)
|
||||
)
|
||||
self._stop_cloudwatch_agent()
|
||||
self._start_cloudwatch_agent(cwa_param_name)
|
||||
|
||||
def _stop_cloudwatch_agent(self) -> None:
|
||||
"""stop Unified CloudWatch Agent"""
|
||||
logger.info(
|
||||
"Stopping Unified CloudWatch Agent package on node {}.".format(self.node_id)
|
||||
)
|
||||
parameters_stop_cwa = {
|
||||
"action": ["stop"],
|
||||
"mode": ["ec2"],
|
||||
}
|
||||
# don't retry failed stop commands
|
||||
# (there's not always an agent to stop)
|
||||
self._ssm_command_waiter(
|
||||
"AmazonCloudWatch-ManageAgent",
|
||||
parameters_stop_cwa,
|
||||
self.node_id,
|
||||
False,
|
||||
)
|
||||
logger.info("Unified CloudWatch Agent stopped on node {}.".format(self.node_id))
|
||||
|
||||
def _start_cloudwatch_agent(self, cwa_param_name: str) -> None:
|
||||
"""start Unified CloudWatch Agent"""
|
||||
logger.info(
|
||||
"Starting Unified CloudWatch Agent package on node {}.".format(self.node_id)
|
||||
)
|
||||
parameters_start_cwa = {
|
||||
"action": ["configure"],
|
||||
"mode": ["ec2"],
|
||||
"optionalConfigurationSource": ["ssm"],
|
||||
"optionalConfigurationLocation": [cwa_param_name],
|
||||
"optionalRestart": ["yes"],
|
||||
}
|
||||
self._ssm_command_waiter(
|
||||
"AmazonCloudWatch-ManageAgent", parameters_start_cwa, self.node_id
|
||||
)
|
||||
logger.info(
|
||||
"Unified CloudWatch Agent started successfully on node {}.".format(
|
||||
self.node_id
|
||||
)
|
||||
)
|
||||
|
||||
def _setup_cwa(self) -> bool:
|
||||
cwa_installed = self._check_cwa_installed_ec2_tag()
|
||||
if cwa_installed == "False":
|
||||
res_cwa_installed = self._ensure_cwa_installed_ssm(self.node_id)
|
||||
return res_cwa_installed
|
||||
else:
|
||||
return True
|
||||
|
||||
def _get_head_node_config_hash(self, config_type: str) -> str:
|
||||
hash_key_value = "-".join([CLOUDWATCH_CONFIG_HASH_TAG_BASE, config_type])
|
||||
filters = copy.deepcopy(
|
||||
self._get_current_cluster_session_nodes(self.cluster_name)
|
||||
)
|
||||
filters.append(
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_NODE_KIND),
|
||||
"Values": [NODE_KIND_HEAD],
|
||||
}
|
||||
)
|
||||
try:
|
||||
instance = list(self.ec2_resource.instances.filter(Filters=filters))
|
||||
assert len(instance) == 1, "More than 1 head node found!"
|
||||
for tag in instance[0].tags:
|
||||
if tag["Key"] == hash_key_value:
|
||||
return tag["Value"]
|
||||
except botocore.exceptions.ClientError as e:
|
||||
logger.warning(
|
||||
"{} Error caught when getting value of {} tag on head node".format(
|
||||
e.response["Error"], hash_key_value
|
||||
)
|
||||
)
|
||||
|
||||
def _get_cur_node_config_hash(self, config_type: str) -> str:
|
||||
hash_key_value = "-".join([CLOUDWATCH_CONFIG_HASH_TAG_BASE, config_type])
|
||||
try:
|
||||
response = self.ec2_client.describe_instances(InstanceIds=[self.node_id])
|
||||
reservations = response["Reservations"]
|
||||
message = "More than 1 response received from describing current node"
|
||||
assert len(reservations) == 1, message
|
||||
instances = reservations[0]["Instances"]
|
||||
assert len(reservations) == 1, message
|
||||
tags = instances[0]["Tags"]
|
||||
hash_value = self._get_default_empty_config_file_hash()
|
||||
for tag in tags:
|
||||
if tag["Key"] == hash_key_value:
|
||||
logger.info(
|
||||
"Successfully get cloudwatch {} hash tag value from "
|
||||
"node {}".format(config_type, self.node_id)
|
||||
)
|
||||
hash_value = tag["Value"]
|
||||
return hash_value
|
||||
except botocore.exceptions.ClientError as e:
|
||||
logger.warning(
|
||||
"{} Error caught when getting hash tag {} tag".format(
|
||||
e.response["Error"], hash_key_value
|
||||
)
|
||||
)
|
||||
|
||||
def _ensure_cwa_installed_ssm(self, node_id: str) -> bool:
|
||||
"""
|
||||
Check if Unified Cloudwatch Agent is installed via ssm run command.
|
||||
If not, notify user to use an AMI with
|
||||
the Unified CloudWatch Agent installed.
|
||||
"""
|
||||
logger.info(
|
||||
"Checking Unified Cloudwatch Agent status on node {}".format(node_id)
|
||||
)
|
||||
parameters_status_cwa = {
|
||||
"action": ["status"],
|
||||
"mode": ["ec2"],
|
||||
}
|
||||
self._ec2_health_check_waiter(node_id)
|
||||
cmd_invocation_res = self._ssm_command_waiter(
|
||||
"AmazonCloudWatch-ManageAgent", parameters_status_cwa, node_id, False
|
||||
)
|
||||
cwa_installed = cmd_invocation_res.get(node_id, False)
|
||||
if not cwa_installed:
|
||||
logger.warning(
|
||||
"Unified CloudWatch Agent not installed on {}. "
|
||||
"Ray logs, metrics not picked up. "
|
||||
"Please use an AMI with Unified CloudWatch Agent installed.".format(
|
||||
node_id
|
||||
)
|
||||
)
|
||||
return False
|
||||
else:
|
||||
return True
|
||||
|
||||
def _get_current_cluster_session_nodes(self, cluster_name: str) -> List[dict]:
|
||||
filters = [
|
||||
{
|
||||
"Name": "instance-state-name",
|
||||
"Values": ["pending", "running"],
|
||||
},
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_CLUSTER_NAME),
|
||||
"Values": [cluster_name],
|
||||
},
|
||||
]
|
||||
return filters
|
||||
|
||||
def _check_cwa_installed_ec2_tag(self) -> List[str]:
|
||||
"""
|
||||
Filtering all nodes to get nodes
|
||||
without Unified CloudWatch Agent installed
|
||||
"""
|
||||
try:
|
||||
response = self.ec2_client.describe_instances(InstanceIds=[self.node_id])
|
||||
reservations = response["Reservations"]
|
||||
message = "More than 1 response received from describing current node"
|
||||
assert len(reservations) == 1, message
|
||||
instances = reservations[0]["Instances"]
|
||||
assert len(instances) == 1, message
|
||||
tags = instances[0]["Tags"]
|
||||
cwa_installed = str(False)
|
||||
for tag in tags:
|
||||
if tag["Key"] == CLOUDWATCH_AGENT_INSTALLED_TAG:
|
||||
logger.info(
|
||||
"Unified CloudWatch Agent is installed on "
|
||||
"node {}".format(self.node_id)
|
||||
)
|
||||
cwa_installed = tag["Value"]
|
||||
return cwa_installed
|
||||
except botocore.exceptions.ClientError as e:
|
||||
logger.warning(
|
||||
"{} Error caught when getting Unified CloudWatch Agent "
|
||||
"status based on {} tag".format(
|
||||
e.response["Error"], CLOUDWATCH_AGENT_INSTALLED_TAG
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def resolve_instance_profile_name(
|
||||
config: Dict[str, Any], default_instance_profile_name: str
|
||||
) -> str:
|
||||
"""Get default cloudwatch instance profile name.
|
||||
|
||||
Args:
|
||||
config: provider section of cluster config file.
|
||||
default_instance_profile_name: default ray instance profile name.
|
||||
|
||||
Returns:
|
||||
default cloudwatch instance profile name if cloudwatch config file
|
||||
exists.
|
||||
default ray instance profile name if cloudwatch config file
|
||||
doesn't exist.
|
||||
"""
|
||||
cwa_cfg_exists = CloudwatchHelper.cloudwatch_config_exists(
|
||||
config, CloudwatchConfigType.AGENT.value
|
||||
)
|
||||
return (
|
||||
CLOUDWATCH_RAY_INSTANCE_PROFILE
|
||||
if cwa_cfg_exists
|
||||
else default_instance_profile_name
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def resolve_iam_role_name(
|
||||
config: Dict[str, Any], default_iam_role_name: str
|
||||
) -> str:
|
||||
"""Get default cloudwatch iam role name.
|
||||
|
||||
Args:
|
||||
config: provider section of cluster config file.
|
||||
default_iam_role_name: default ray iam role name.
|
||||
|
||||
Returns:
|
||||
default cloudwatch iam role name if cloudwatch config file exists.
|
||||
default ray iam role name if cloudwatch config file doesn't exist.
|
||||
"""
|
||||
cwa_cfg_exists = CloudwatchHelper.cloudwatch_config_exists(
|
||||
config, CloudwatchConfigType.AGENT.value
|
||||
)
|
||||
return CLOUDWATCH_RAY_IAM_ROLE if cwa_cfg_exists else default_iam_role_name
|
||||
|
||||
@staticmethod
|
||||
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, CloudwatchConfigType.AGENT.value
|
||||
)
|
||||
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
|
||||
|
||||
@staticmethod
|
||||
def cloudwatch_config_exists(config: Dict[str, Any], config_type: str) -> bool:
|
||||
"""Check if CloudWatch configuration was specified by the user
|
||||
in their cluster config file.
|
||||
|
||||
Specifically, this function checks if a CloudWatch config file is
|
||||
specified by the user in their cluster config file.
|
||||
|
||||
Args:
|
||||
config: provider section of cluster config file.
|
||||
config_type: type of CloudWatch config file.
|
||||
|
||||
Returns:
|
||||
True if config file is specified by user.
|
||||
False if config file is not specified.
|
||||
"""
|
||||
cfg = config.get("cloudwatch", {}).get(config_type, {}).get("config")
|
||||
return bool(cfg)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,715 @@
|
||||
import copy
|
||||
import logging
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections import OrderedDict, defaultdict
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import botocore
|
||||
from boto3.resources.base import ServiceResource
|
||||
|
||||
import ray
|
||||
import ray._private.ray_constants as ray_constants
|
||||
from ray.autoscaler._private.aws.cloudwatch.cloudwatch_helper import (
|
||||
CLOUDWATCH_AGENT_INSTALLED_AMI_TAG,
|
||||
CLOUDWATCH_AGENT_INSTALLED_TAG,
|
||||
CloudwatchHelper,
|
||||
)
|
||||
from ray.autoscaler._private.aws.config import bootstrap_aws
|
||||
from ray.autoscaler._private.aws.utils import (
|
||||
boto_exception_handler,
|
||||
client_cache,
|
||||
resource_cache,
|
||||
)
|
||||
from ray.autoscaler._private.cli_logger import cf, cli_logger
|
||||
from ray.autoscaler._private.constants import BOTO_CREATE_MAX_RETRIES, BOTO_MAX_RETRIES
|
||||
from ray.autoscaler._private.log_timer import LogTimer
|
||||
from ray.autoscaler.node_launch_exception import NodeLaunchException
|
||||
from ray.autoscaler.node_provider import NodeProvider
|
||||
from ray.autoscaler.tags import (
|
||||
TAG_RAY_CLUSTER_NAME,
|
||||
TAG_RAY_LAUNCH_CONFIG,
|
||||
TAG_RAY_NODE_KIND,
|
||||
TAG_RAY_NODE_NAME,
|
||||
TAG_RAY_USER_NODE_TYPE,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
TAG_BATCH_DELAY = 1
|
||||
LIST_RETRY_DELAY_SEC = 1
|
||||
|
||||
|
||||
def to_aws_format(tags):
|
||||
"""Convert the Ray node name tag to the AWS-specific 'Name' tag."""
|
||||
|
||||
if TAG_RAY_NODE_NAME in tags:
|
||||
tags["Name"] = tags[TAG_RAY_NODE_NAME]
|
||||
del tags[TAG_RAY_NODE_NAME]
|
||||
return tags
|
||||
|
||||
|
||||
def from_aws_format(tags):
|
||||
"""Convert the AWS-specific 'Name' tag to the Ray node name tag."""
|
||||
|
||||
if "Name" in tags:
|
||||
tags[TAG_RAY_NODE_NAME] = tags["Name"]
|
||||
del tags["Name"]
|
||||
return tags
|
||||
|
||||
|
||||
def make_ec2_resource(region, max_retries, aws_credentials=None) -> ServiceResource:
|
||||
"""Make client, retrying requests up to `max_retries`."""
|
||||
aws_credentials = aws_credentials or {}
|
||||
return resource_cache("ec2", region, max_retries, **aws_credentials)
|
||||
|
||||
|
||||
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: the region of the AWS provider. e.g., "us-west-2".
|
||||
aws_credentials: AWS credentials to use for the boto3 client.
|
||||
|
||||
Returns:
|
||||
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
|
||||
|
||||
|
||||
class AWSNodeProvider(NodeProvider):
|
||||
max_terminate_nodes = 1000
|
||||
|
||||
def __init__(self, provider_config, cluster_name):
|
||||
NodeProvider.__init__(self, provider_config, cluster_name)
|
||||
self.cache_stopped_nodes = provider_config.get("cache_stopped_nodes", True)
|
||||
aws_credentials = provider_config.get("aws_credentials")
|
||||
|
||||
self.ec2 = make_ec2_resource(
|
||||
region=provider_config["region"],
|
||||
max_retries=BOTO_MAX_RETRIES,
|
||||
aws_credentials=aws_credentials,
|
||||
)
|
||||
self.ec2_fail_fast = make_ec2_resource(
|
||||
region=provider_config["region"],
|
||||
max_retries=0,
|
||||
aws_credentials=aws_credentials,
|
||||
)
|
||||
|
||||
# Tags that we believe to actually be on EC2.
|
||||
self.tag_cache = {}
|
||||
# Tags that we will soon upload.
|
||||
self.tag_cache_pending = defaultdict(dict)
|
||||
# Number of threads waiting for a batched tag update.
|
||||
self.batch_thread_count = 0
|
||||
self.batch_update_done = threading.Event()
|
||||
self.batch_update_done.set()
|
||||
self.ready_for_new_batch = threading.Event()
|
||||
self.ready_for_new_batch.set()
|
||||
self.tag_cache_lock = threading.Lock()
|
||||
self.count_lock = threading.Lock()
|
||||
# Prevent concurrent create_node calls to get the same stopped/stopping node to reuse.
|
||||
self._reuse_node_lock = threading.Lock()
|
||||
|
||||
# Cache of node objects from the last nodes() call. This avoids
|
||||
# excessive DescribeInstances requests.
|
||||
self.cached_nodes = {}
|
||||
|
||||
def non_terminated_nodes(self, tag_filters):
|
||||
# Note that these filters are acceptable because they are set on
|
||||
# node initialization, and so can never be sitting in the cache.
|
||||
tag_filters = to_aws_format(tag_filters)
|
||||
filters = [
|
||||
{
|
||||
"Name": "instance-state-name",
|
||||
"Values": ["pending", "running"],
|
||||
},
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_CLUSTER_NAME),
|
||||
"Values": [self.cluster_name],
|
||||
},
|
||||
]
|
||||
for k, v in tag_filters.items():
|
||||
filters.append(
|
||||
{
|
||||
"Name": "tag:{}".format(k),
|
||||
"Values": [v],
|
||||
}
|
||||
)
|
||||
|
||||
with boto_exception_handler("Failed to fetch running instances from AWS."):
|
||||
nodes = list(self.ec2.instances.filter(Filters=filters))
|
||||
|
||||
# Populate the tag cache with initial information if necessary
|
||||
for node in nodes:
|
||||
if node.id in self.tag_cache:
|
||||
continue
|
||||
|
||||
self.tag_cache[node.id] = from_aws_format(
|
||||
{x["Key"]: x["Value"] for x in node.tags}
|
||||
)
|
||||
|
||||
self.cached_nodes = {node.id: node for node in nodes}
|
||||
return [node.id for node in nodes]
|
||||
|
||||
def is_running(self, node_id):
|
||||
node = self._get_cached_node(node_id)
|
||||
return node.state["Name"] == "running"
|
||||
|
||||
def is_terminated(self, node_id):
|
||||
node = self._get_cached_node(node_id)
|
||||
state = node.state["Name"]
|
||||
return state not in ["running", "pending"]
|
||||
|
||||
def node_tags(self, node_id):
|
||||
with self.tag_cache_lock:
|
||||
d1 = self.tag_cache[node_id]
|
||||
d2 = self.tag_cache_pending.get(node_id, {})
|
||||
return dict(d1, **d2)
|
||||
|
||||
def external_ip(self, node_id):
|
||||
node = self._get_cached_node(node_id)
|
||||
|
||||
if node.public_ip_address is None:
|
||||
node = self._get_node(node_id)
|
||||
|
||||
return node.public_ip_address
|
||||
|
||||
def internal_ip(self, node_id):
|
||||
node = self._get_cached_node(node_id)
|
||||
|
||||
if node.private_ip_address is None:
|
||||
node = self._get_node(node_id)
|
||||
|
||||
return node.private_ip_address
|
||||
|
||||
def set_node_tags(self, node_id, tags):
|
||||
is_batching_thread = False
|
||||
with self.tag_cache_lock:
|
||||
if not self.tag_cache_pending:
|
||||
is_batching_thread = True
|
||||
# Wait for threads in the last batch to exit
|
||||
self.ready_for_new_batch.wait()
|
||||
self.ready_for_new_batch.clear()
|
||||
self.batch_update_done.clear()
|
||||
self.tag_cache_pending[node_id].update(tags)
|
||||
|
||||
if is_batching_thread:
|
||||
time.sleep(TAG_BATCH_DELAY)
|
||||
with self.tag_cache_lock:
|
||||
self._update_node_tags()
|
||||
self.batch_update_done.set()
|
||||
|
||||
with self.count_lock:
|
||||
self.batch_thread_count += 1
|
||||
self.batch_update_done.wait()
|
||||
|
||||
with self.count_lock:
|
||||
self.batch_thread_count -= 1
|
||||
if self.batch_thread_count == 0:
|
||||
self.ready_for_new_batch.set()
|
||||
|
||||
def _update_node_tags(self):
|
||||
batch_updates = defaultdict(list)
|
||||
|
||||
for node_id, tags in self.tag_cache_pending.items():
|
||||
for x in tags.items():
|
||||
batch_updates[x].append(node_id)
|
||||
self.tag_cache[node_id].update(tags)
|
||||
|
||||
self.tag_cache_pending = defaultdict(dict)
|
||||
|
||||
self._create_tags(batch_updates)
|
||||
|
||||
def _create_tags(self, batch_updates):
|
||||
for (k, v), node_ids in batch_updates.items():
|
||||
m = "Set tag {}={} on {}".format(k, v, node_ids)
|
||||
with LogTimer("AWSNodeProvider: {}".format(m)):
|
||||
if k == TAG_RAY_NODE_NAME:
|
||||
k = "Name"
|
||||
self.ec2.meta.client.create_tags(
|
||||
Resources=node_ids,
|
||||
Tags=[{"Key": k, "Value": v}],
|
||||
)
|
||||
|
||||
def create_node(self, node_config, tags, count) -> Dict[str, Any]:
|
||||
"""Creates instances.
|
||||
|
||||
Returns dict mapping instance id to ec2.Instance object for the created
|
||||
instances.
|
||||
"""
|
||||
# sort tags by key to support deterministic unit test stubbing
|
||||
tags = OrderedDict(sorted(copy.deepcopy(tags).items()))
|
||||
|
||||
reused_nodes_dict = {}
|
||||
# Try to reuse previously stopped nodes with compatible configs
|
||||
if self.cache_stopped_nodes:
|
||||
# TODO(ekl) this is breaking the abstraction boundary a little by
|
||||
# peeking into the tag set.
|
||||
filters = [
|
||||
{
|
||||
"Name": "instance-state-name",
|
||||
"Values": ["stopped", "stopping"],
|
||||
},
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_CLUSTER_NAME),
|
||||
"Values": [self.cluster_name],
|
||||
},
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_NODE_KIND),
|
||||
"Values": [tags[TAG_RAY_NODE_KIND]],
|
||||
},
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_LAUNCH_CONFIG),
|
||||
"Values": [tags[TAG_RAY_LAUNCH_CONFIG]],
|
||||
},
|
||||
]
|
||||
# This tag may not always be present.
|
||||
if TAG_RAY_USER_NODE_TYPE in tags:
|
||||
filters.append(
|
||||
{
|
||||
"Name": "tag:{}".format(TAG_RAY_USER_NODE_TYPE),
|
||||
"Values": [tags[TAG_RAY_USER_NODE_TYPE]],
|
||||
}
|
||||
)
|
||||
|
||||
with self._reuse_node_lock:
|
||||
reuse_nodes = list(self.ec2.instances.filter(Filters=filters))[:count]
|
||||
reuse_node_ids = [n.id for n in reuse_nodes]
|
||||
reused_nodes_dict = {n.id: n for n in reuse_nodes}
|
||||
if reuse_nodes:
|
||||
cli_logger.print(
|
||||
# todo: handle plural vs singular?
|
||||
"Reusing nodes {}. "
|
||||
"To disable reuse, set `cache_stopped_nodes: False` "
|
||||
"under `provider` in the cluster configuration.",
|
||||
cli_logger.render_list(reuse_node_ids),
|
||||
)
|
||||
|
||||
# todo: timed?
|
||||
with cli_logger.group("Stopping instances to reuse"):
|
||||
for node in reuse_nodes:
|
||||
self.tag_cache[node.id] = from_aws_format(
|
||||
{x["Key"]: x["Value"] for x in node.tags}
|
||||
)
|
||||
if node.state["Name"] == "stopping":
|
||||
cli_logger.print(
|
||||
"Waiting for instance {} to stop", node.id
|
||||
)
|
||||
node.wait_until_stopped()
|
||||
|
||||
self.ec2.meta.client.start_instances(InstanceIds=reuse_node_ids)
|
||||
for node_id in reuse_node_ids:
|
||||
self.set_node_tags(node_id, tags)
|
||||
count -= len(reuse_node_ids)
|
||||
|
||||
created_nodes_dict = {}
|
||||
if count:
|
||||
created_nodes_dict = self._create_node(node_config, tags, count)
|
||||
|
||||
all_created_nodes = reused_nodes_dict
|
||||
all_created_nodes.update(created_nodes_dict)
|
||||
return all_created_nodes
|
||||
|
||||
@staticmethod
|
||||
def _merge_tag_specs(
|
||||
tag_specs: List[Dict[str, Any]], user_tag_specs: List[Dict[str, Any]]
|
||||
) -> None:
|
||||
"""
|
||||
Merges user-provided node config tag specifications into a base
|
||||
list of node provider tag specifications. The base list of
|
||||
node provider tag specs is modified in-place.
|
||||
|
||||
This allows users to add tags and override values of existing
|
||||
tags with their own, and only applies to the resource type
|
||||
"instance". All other resource types are appended to the list of
|
||||
tag specs.
|
||||
|
||||
Args:
|
||||
tag_specs: base node provider tag specs
|
||||
user_tag_specs: user's node config tag specs
|
||||
"""
|
||||
|
||||
for user_tag_spec in user_tag_specs:
|
||||
if user_tag_spec["ResourceType"] == "instance":
|
||||
for user_tag in user_tag_spec["Tags"]:
|
||||
exists = False
|
||||
for tag in tag_specs[0]["Tags"]:
|
||||
if user_tag["Key"] == tag["Key"]:
|
||||
exists = True
|
||||
tag["Value"] = user_tag["Value"]
|
||||
break
|
||||
if not exists:
|
||||
tag_specs[0]["Tags"] += [user_tag]
|
||||
else:
|
||||
tag_specs += [user_tag_spec]
|
||||
|
||||
def _create_node(self, node_config, tags, count):
|
||||
created_nodes_dict = {}
|
||||
|
||||
tags = to_aws_format(tags)
|
||||
conf = node_config.copy()
|
||||
|
||||
tag_pairs = [
|
||||
{
|
||||
"Key": TAG_RAY_CLUSTER_NAME,
|
||||
"Value": self.cluster_name,
|
||||
}
|
||||
]
|
||||
for k, v in tags.items():
|
||||
tag_pairs.append(
|
||||
{
|
||||
"Key": k,
|
||||
"Value": v,
|
||||
}
|
||||
)
|
||||
if CloudwatchHelper.cloudwatch_config_exists(self.provider_config, "agent"):
|
||||
cwa_installed = self._check_ami_cwa_installation(node_config)
|
||||
if cwa_installed:
|
||||
tag_pairs.extend(
|
||||
[
|
||||
{
|
||||
"Key": CLOUDWATCH_AGENT_INSTALLED_TAG,
|
||||
"Value": "True",
|
||||
}
|
||||
]
|
||||
)
|
||||
tag_specs = [
|
||||
{
|
||||
"ResourceType": "instance",
|
||||
"Tags": tag_pairs,
|
||||
}
|
||||
]
|
||||
user_tag_specs = conf.get("TagSpecifications", [])
|
||||
AWSNodeProvider._merge_tag_specs(tag_specs, user_tag_specs)
|
||||
|
||||
# SubnetIds is not a real config key: we must resolve to a
|
||||
# single SubnetId before invoking the AWS API.
|
||||
subnet_ids = conf.pop("SubnetIds")
|
||||
|
||||
# update config with min/max node counts and tag specs
|
||||
conf.update({"MinCount": 1, "MaxCount": count, "TagSpecifications": tag_specs})
|
||||
|
||||
# Try to always launch in the first listed subnet.
|
||||
subnet_idx = 0
|
||||
cli_logger_tags = {}
|
||||
# NOTE: This ensures that we try ALL availability zones before
|
||||
# throwing an error.
|
||||
max_tries = max(BOTO_CREATE_MAX_RETRIES, len(subnet_ids))
|
||||
for attempt in range(1, max_tries + 1):
|
||||
try:
|
||||
if "NetworkInterfaces" in conf:
|
||||
net_ifs = conf["NetworkInterfaces"]
|
||||
# remove security group IDs previously copied from network
|
||||
# interfaces (create_instances call fails otherwise)
|
||||
conf.pop("SecurityGroupIds", None)
|
||||
cli_logger_tags["network_interfaces"] = str(net_ifs)
|
||||
else:
|
||||
subnet_id = subnet_ids[subnet_idx % len(subnet_ids)]
|
||||
conf["SubnetId"] = subnet_id
|
||||
cli_logger_tags["subnet_id"] = subnet_id
|
||||
|
||||
created = self.ec2_fail_fast.create_instances(**conf)
|
||||
created_nodes_dict = {n.id: n for n in created}
|
||||
|
||||
# todo: timed?
|
||||
# todo: handle plurality?
|
||||
with cli_logger.group(
|
||||
"Launched {} nodes", count, _tags=cli_logger_tags
|
||||
):
|
||||
for instance in created:
|
||||
# NOTE(maximsmol): This is needed for mocking
|
||||
# boto3 for tests. This is likely a bug in moto
|
||||
# but AWS docs don't seem to say.
|
||||
# You can patch moto/ec2/responses/instances.py
|
||||
# to fix this (add <stateReason> to EC2_RUN_INSTANCES)
|
||||
|
||||
# The correct value is technically
|
||||
# {"code": "0", "Message": "pending"}
|
||||
state_reason = "pending"
|
||||
if instance.state_reason:
|
||||
state_reason = (
|
||||
instance.state_reason["Message"] or state_reason
|
||||
)
|
||||
|
||||
cli_logger.print(
|
||||
"Launched instance {}",
|
||||
instance.instance_id,
|
||||
_tags=dict(
|
||||
state=instance.state["Name"],
|
||||
info=state_reason,
|
||||
),
|
||||
)
|
||||
break
|
||||
except botocore.exceptions.ClientError as exc:
|
||||
# Launch failure may be due to instance type availability in
|
||||
# the given AZ
|
||||
subnet_idx += 1
|
||||
if attempt == max_tries:
|
||||
try:
|
||||
exc = NodeLaunchException(
|
||||
category=exc.response["Error"]["Code"],
|
||||
description=exc.response["Error"]["Message"],
|
||||
src_exc_info=sys.exc_info(),
|
||||
)
|
||||
except Exception:
|
||||
# In theory, all ClientError's we expect to get should
|
||||
# have these fields, but just in case we can't parse
|
||||
# it, it's fine, just throw the original error.
|
||||
logger.warning("Couldn't parse exception.", exc)
|
||||
pass
|
||||
cli_logger.abort(
|
||||
"Failed to launch instances. Max attempts exceeded.",
|
||||
exc=exc,
|
||||
)
|
||||
else:
|
||||
cli_logger.warning(
|
||||
"create_instances: Attempt failed with {}, retrying.", exc
|
||||
)
|
||||
|
||||
return created_nodes_dict
|
||||
|
||||
def terminate_node(self, node_id):
|
||||
node = self._get_cached_node(node_id)
|
||||
if self.cache_stopped_nodes:
|
||||
if node.spot_instance_request_id:
|
||||
cli_logger.print(
|
||||
"Terminating instance {} "
|
||||
+ cf.dimmed("(cannot stop spot instances, only terminate)"),
|
||||
node_id,
|
||||
) # todo: show node name?
|
||||
node.terminate()
|
||||
else:
|
||||
cli_logger.print(
|
||||
"Stopping instance {} "
|
||||
+ cf.dimmed(
|
||||
"(to terminate instead, "
|
||||
"set `cache_stopped_nodes: False` "
|
||||
"under `provider` in the cluster configuration)"
|
||||
),
|
||||
node_id,
|
||||
) # todo: show node name?
|
||||
node.stop()
|
||||
else:
|
||||
node.terminate()
|
||||
|
||||
# TODO (Alex): We are leaking the tag cache here. Naively, we would
|
||||
# want to just remove the cache entry here, but terminating can be
|
||||
# asyncrhonous or error, which would result in a use after free error.
|
||||
# If this leak becomes bad, we can garbage collect the tag cache when
|
||||
# the node cache is updated.
|
||||
|
||||
def _check_ami_cwa_installation(self, config):
|
||||
response = self.ec2.meta.client.describe_images(ImageIds=[config["ImageId"]])
|
||||
cwa_installed = False
|
||||
images = response.get("Images")
|
||||
if images:
|
||||
assert len(images) == 1, (
|
||||
f"Expected to find only 1 AMI with the given ID, "
|
||||
f"but found {len(images)}."
|
||||
)
|
||||
image_name = images[0].get("Name", "")
|
||||
if CLOUDWATCH_AGENT_INSTALLED_AMI_TAG in image_name:
|
||||
cwa_installed = True
|
||||
return cwa_installed
|
||||
|
||||
def terminate_nodes(self, node_ids):
|
||||
if not node_ids:
|
||||
return
|
||||
|
||||
terminate_instances_func = self.ec2.meta.client.terminate_instances
|
||||
stop_instances_func = self.ec2.meta.client.stop_instances
|
||||
|
||||
# In some cases, this function stops some nodes, but terminates others.
|
||||
# Each of these requires a different EC2 API call. So, we use the
|
||||
# "nodes_to_terminate" dict below to keep track of exactly which API
|
||||
# call will be used to stop/terminate which set of nodes. The key is
|
||||
# the function to use, and the value is the list of nodes to terminate
|
||||
# with that function.
|
||||
nodes_to_terminate = {terminate_instances_func: [], stop_instances_func: []}
|
||||
|
||||
if self.cache_stopped_nodes:
|
||||
spot_ids = []
|
||||
on_demand_ids = []
|
||||
|
||||
for node_id in node_ids:
|
||||
if self._get_cached_node(node_id).spot_instance_request_id:
|
||||
spot_ids += [node_id]
|
||||
else:
|
||||
on_demand_ids += [node_id]
|
||||
|
||||
if on_demand_ids:
|
||||
# todo: show node names?
|
||||
cli_logger.print(
|
||||
"Stopping instances {} "
|
||||
+ cf.dimmed(
|
||||
"(to terminate instead, "
|
||||
"set `cache_stopped_nodes: False` "
|
||||
"under `provider` in the cluster configuration)"
|
||||
),
|
||||
cli_logger.render_list(on_demand_ids),
|
||||
)
|
||||
|
||||
if spot_ids:
|
||||
cli_logger.print(
|
||||
"Terminating instances {} "
|
||||
+ cf.dimmed("(cannot stop spot instances, only terminate)"),
|
||||
cli_logger.render_list(spot_ids),
|
||||
)
|
||||
|
||||
nodes_to_terminate[stop_instances_func] = on_demand_ids
|
||||
nodes_to_terminate[terminate_instances_func] = spot_ids
|
||||
else:
|
||||
nodes_to_terminate[terminate_instances_func] = node_ids
|
||||
|
||||
max_terminate_nodes = (
|
||||
self.max_terminate_nodes
|
||||
if self.max_terminate_nodes is not None
|
||||
else len(node_ids)
|
||||
)
|
||||
|
||||
for terminate_func, nodes in nodes_to_terminate.items():
|
||||
for start in range(0, len(nodes), max_terminate_nodes):
|
||||
terminate_func(InstanceIds=nodes[start : start + max_terminate_nodes])
|
||||
|
||||
def _get_node(self, node_id):
|
||||
"""Refresh and get info for this node, updating the cache."""
|
||||
self.non_terminated_nodes({}) # Side effect: updates cache
|
||||
|
||||
if node_id in self.cached_nodes:
|
||||
return self.cached_nodes[node_id]
|
||||
|
||||
# Node not in {pending, running} -- retry with a point query. This
|
||||
# usually means the node was recently preempted or terminated.
|
||||
# The EC2 API is eventually consistent. This means that an instance
|
||||
# might not be immediately visible. So we need to retry the query a few times.
|
||||
# See: https://docs.aws.amazon.com/ec2/latest/devguide/eventual-consistency.html
|
||||
# and https://github.com/ray-project/ray/issues/51861
|
||||
for attempts in range(max(BOTO_MAX_RETRIES, 1)): # at least try once.
|
||||
matches = list(self.ec2.instances.filter(InstanceIds=[node_id]))
|
||||
if len(matches) == 1:
|
||||
return matches[0]
|
||||
cli_logger.warning(
|
||||
"Attempt to fetch EC2 instances that have instance ID {}. Got {} matching EC2 instances. Will retry after {} second. This is retry number {}, and the maximum number of retries is {}.",
|
||||
node_id,
|
||||
len(matches),
|
||||
LIST_RETRY_DELAY_SEC,
|
||||
attempts + 1,
|
||||
BOTO_MAX_RETRIES,
|
||||
)
|
||||
time.sleep(LIST_RETRY_DELAY_SEC)
|
||||
raise AssertionError("Invalid instance id {}".format(node_id))
|
||||
|
||||
def _get_cached_node(self, node_id):
|
||||
"""Return node info from cache if possible, otherwise fetches it."""
|
||||
if node_id in self.cached_nodes:
|
||||
return self.cached_nodes[node_id]
|
||||
|
||||
return self._get_node(node_id)
|
||||
|
||||
@staticmethod
|
||||
def bootstrap_config(cluster_config):
|
||||
return bootstrap_aws(cluster_config)
|
||||
|
||||
@staticmethod
|
||||
def fillout_available_node_types_resources(
|
||||
cluster_config: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Fills out missing "resources" field for available_node_types."""
|
||||
if "available_node_types" not in cluster_config:
|
||||
return cluster_config
|
||||
cluster_config = copy.deepcopy(cluster_config)
|
||||
|
||||
instances_list = list_ec2_instances(
|
||||
cluster_config["provider"]["region"],
|
||||
cluster_config["provider"].get("aws_credentials"),
|
||||
)
|
||||
instances_dict = {
|
||||
instance["InstanceType"]: instance for instance in instances_list
|
||||
}
|
||||
available_node_types = cluster_config["available_node_types"]
|
||||
head_node_type = cluster_config["head_node_type"]
|
||||
for node_type in available_node_types:
|
||||
instance_type = available_node_types[node_type]["node_config"][
|
||||
"InstanceType"
|
||||
]
|
||||
if instance_type in instances_dict:
|
||||
cpus = instances_dict[instance_type]["VCpuInfo"]["DefaultVCpus"]
|
||||
|
||||
autodetected_resources = {"CPU": cpus}
|
||||
if node_type != head_node_type:
|
||||
# we only autodetect worker node type memory resource
|
||||
memory_total = instances_dict[instance_type]["MemoryInfo"][
|
||||
"SizeInMiB"
|
||||
]
|
||||
memory_total = int(memory_total) * 1024 * 1024
|
||||
prop = 1 - ray_constants.DEFAULT_OBJECT_STORE_MEMORY_PROPORTION
|
||||
memory_resources = int(memory_total * prop)
|
||||
autodetected_resources["memory"] = memory_resources
|
||||
|
||||
for (
|
||||
accelerator_manager
|
||||
) in ray._private.accelerators.get_all_accelerator_managers():
|
||||
num_accelerators = (
|
||||
accelerator_manager.get_ec2_instance_num_accelerators(
|
||||
instance_type, instances_dict
|
||||
)
|
||||
)
|
||||
accelerator_type = (
|
||||
accelerator_manager.get_ec2_instance_accelerator_type(
|
||||
instance_type, instances_dict
|
||||
)
|
||||
)
|
||||
if num_accelerators:
|
||||
autodetected_resources[
|
||||
accelerator_manager.get_resource_name()
|
||||
] = num_accelerators
|
||||
if accelerator_type:
|
||||
autodetected_resources[
|
||||
f"accelerator_type:{accelerator_type}"
|
||||
] = 1
|
||||
|
||||
autodetected_resources.update(
|
||||
available_node_types[node_type].get("resources", {})
|
||||
)
|
||||
if autodetected_resources != available_node_types[node_type].get(
|
||||
"resources", {}
|
||||
):
|
||||
available_node_types[node_type][
|
||||
"resources"
|
||||
] = autodetected_resources
|
||||
logger.debug(
|
||||
"Updating the resources of {} to {}.".format(
|
||||
node_type, autodetected_resources
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Instance type "
|
||||
+ instance_type
|
||||
+ " is not available in AWS region: "
|
||||
+ cluster_config["provider"]["region"]
|
||||
+ "."
|
||||
)
|
||||
return cluster_config
|
||||
@@ -0,0 +1,181 @@
|
||||
from collections import defaultdict
|
||||
from functools import lru_cache
|
||||
|
||||
import boto3
|
||||
from boto3.exceptions import ResourceNotExistsError
|
||||
from boto3.resources.base import ServiceResource
|
||||
from botocore.client import BaseClient
|
||||
from botocore.config import Config
|
||||
|
||||
from ray.autoscaler._private.cli_logger import cf, cli_logger
|
||||
from ray.autoscaler._private.constants import BOTO_MAX_RETRIES
|
||||
|
||||
|
||||
class LazyDefaultDict(defaultdict):
|
||||
"""
|
||||
LazyDefaultDict(default_factory[, ...]) --> dict with default factory
|
||||
|
||||
The default factory is call with the key argument to produce
|
||||
a new value when a key is not present, in __getitem__ only.
|
||||
A LazyDefaultDict compares equal to a dict with the same items.
|
||||
All remaining arguments are treated the same as if they were
|
||||
passed to the dict constructor, including keyword arguments.
|
||||
"""
|
||||
|
||||
def __missing__(self, key):
|
||||
"""
|
||||
__missing__(key) # Called by __getitem__ for missing key; pseudo-code:
|
||||
if self.default_factory is None: raise KeyError((key,))
|
||||
self[key] = value = self.default_factory(key)
|
||||
return value
|
||||
"""
|
||||
self[key] = self.default_factory(key)
|
||||
return self[key]
|
||||
|
||||
|
||||
def handle_boto_error(exc, msg, *args, **kwargs):
|
||||
error_code = None
|
||||
error_info = None
|
||||
# todo: not sure if these exceptions always have response
|
||||
if hasattr(exc, "response"):
|
||||
error_info = exc.response.get("Error", None)
|
||||
if error_info is not None:
|
||||
error_code = error_info.get("Code", None)
|
||||
|
||||
generic_message_args = [
|
||||
"{}\nError code: {}",
|
||||
msg.format(*args, **kwargs),
|
||||
cf.bold(error_code),
|
||||
]
|
||||
|
||||
# apparently
|
||||
# ExpiredTokenException
|
||||
# ExpiredToken
|
||||
# RequestExpired
|
||||
# are all the same pretty much
|
||||
credentials_expiration_codes = [
|
||||
"ExpiredTokenException",
|
||||
"ExpiredToken",
|
||||
"RequestExpired",
|
||||
]
|
||||
|
||||
if error_code in credentials_expiration_codes:
|
||||
# "An error occurred (ExpiredToken) when calling the
|
||||
# GetInstanceProfile operation: The security token
|
||||
# included in the request is expired"
|
||||
|
||||
# "An error occurred (RequestExpired) when calling the
|
||||
# DescribeKeyPairs operation: Request has expired."
|
||||
|
||||
token_command = (
|
||||
"aws sts get-session-token "
|
||||
"--serial-number arn:aws:iam::"
|
||||
+ cf.underlined("ROOT_ACCOUNT_ID")
|
||||
+ ":mfa/"
|
||||
+ cf.underlined("AWS_USERNAME")
|
||||
+ " --token-code "
|
||||
+ cf.underlined("TWO_FACTOR_AUTH_CODE")
|
||||
)
|
||||
|
||||
secret_key_var = (
|
||||
"export AWS_SECRET_ACCESS_KEY = "
|
||||
+ cf.underlined("REPLACE_ME")
|
||||
+ " # found at Credentials.SecretAccessKey"
|
||||
)
|
||||
session_token_var = (
|
||||
"export AWS_SESSION_TOKEN = "
|
||||
+ cf.underlined("REPLACE_ME")
|
||||
+ " # found at Credentials.SessionToken"
|
||||
)
|
||||
access_key_id_var = (
|
||||
"export AWS_ACCESS_KEY_ID = "
|
||||
+ cf.underlined("REPLACE_ME")
|
||||
+ " # found at Credentials.AccessKeyId"
|
||||
)
|
||||
|
||||
# fixme: replace with a Github URL that points
|
||||
# to our repo
|
||||
aws_session_script_url = (
|
||||
"https://gist.github.com/maximsmol/a0284e1d97b25d417bd9ae02e5f450cf"
|
||||
)
|
||||
|
||||
cli_logger.verbose_error(*generic_message_args)
|
||||
cli_logger.verbose(vars(exc))
|
||||
|
||||
cli_logger.panic("Your AWS session has expired.")
|
||||
cli_logger.newline()
|
||||
cli_logger.panic("You can request a new one using")
|
||||
cli_logger.panic(cf.bold(token_command))
|
||||
cli_logger.panic("then expose it to Ray by setting")
|
||||
cli_logger.panic(cf.bold(secret_key_var))
|
||||
cli_logger.panic(cf.bold(session_token_var))
|
||||
cli_logger.panic(cf.bold(access_key_id_var))
|
||||
cli_logger.newline()
|
||||
cli_logger.panic("You can find a script that automates this at:")
|
||||
cli_logger.panic(cf.underlined(aws_session_script_url))
|
||||
# Do not re-raise the exception here because it looks awful
|
||||
# and we already print all the info in verbose
|
||||
cli_logger.abort()
|
||||
|
||||
# todo: any other errors that we should catch separately?
|
||||
|
||||
cli_logger.panic(*generic_message_args)
|
||||
cli_logger.newline()
|
||||
with cli_logger.verbatim_error_ctx("Boto3 error:"):
|
||||
cli_logger.verbose("{}", str(vars(exc)))
|
||||
cli_logger.panic("{}", str(exc))
|
||||
cli_logger.abort()
|
||||
|
||||
|
||||
def boto_exception_handler(msg, *args, **kwargs):
|
||||
# todo: implement timer
|
||||
class ExceptionHandlerContextManager:
|
||||
def __enter__(self):
|
||||
pass
|
||||
|
||||
def __exit__(self, type, value, tb):
|
||||
import botocore
|
||||
|
||||
if type is botocore.exceptions.ClientError:
|
||||
handle_boto_error(value, msg, *args, **kwargs)
|
||||
|
||||
return ExceptionHandlerContextManager()
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def resource_cache(
|
||||
name, region, max_retries=BOTO_MAX_RETRIES, **kwargs
|
||||
) -> ServiceResource:
|
||||
cli_logger.verbose(
|
||||
"Creating AWS resource `{}` in `{}`", cf.bold(name), cf.bold(region)
|
||||
)
|
||||
kwargs.setdefault(
|
||||
"config",
|
||||
Config(retries={"max_attempts": max_retries}),
|
||||
)
|
||||
return boto3.resource(
|
||||
name,
|
||||
region,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def client_cache(name, region, max_retries=BOTO_MAX_RETRIES, **kwargs) -> BaseClient:
|
||||
try:
|
||||
# try to re-use a client from the resource cache first
|
||||
return resource_cache(name, region, max_retries, **kwargs).meta.client
|
||||
except ResourceNotExistsError:
|
||||
# fall back for clients without an associated resource
|
||||
cli_logger.verbose(
|
||||
"Creating AWS client `{}` in `{}`", cf.bold(name), cf.bold(region)
|
||||
)
|
||||
kwargs.setdefault(
|
||||
"config",
|
||||
Config(retries={"max_attempts": max_retries}),
|
||||
)
|
||||
return boto3.client(
|
||||
name,
|
||||
region,
|
||||
**kwargs,
|
||||
)
|
||||
Reference in New Issue
Block a user