wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-07 11:27:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

[autoscaler] Split autoscaler interface public private (#10898)

Browse Source
This commit is contained in:
Eric Liang
2020-09-18 18:16:23 -07:00
committed by GitHub
parent 9a07c7b963
commit 6a227ae501
56 changed files with 436 additions and 251 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/autoscaler/aws/config.py
-731
View File
@@ -1,731 +0,0 @@
from distutils.version import StrictVersion
from functools import lru_cache
from functools import partial
import itertools
import json
import os
import time
import logging
import boto3
from botocore.config import Config
import botocore
from ray.ray_constants import BOTO_MAX_RETRIES
from ray.autoscaler.tags import NODE_KIND_WORKER, NODE_KIND_HEAD
from ray.autoscaler.aws.utils import LazyDefaultDict, handle_boto_error
from ray.autoscaler.node_provider import PROVIDER_PRETTY_NAMES
from ray.autoscaler.cli_logger import cli_logger
import colorful as cf
logger = logging.getLogger(__name__)
RAY = "ray-autoscaler"
DEFAULT_RAY_INSTANCE_PROFILE = RAY + "-v1"
DEFAULT_RAY_IAM_ROLE = RAY + "-v1"
SECURITY_GROUP_TEMPLATE = RAY + "-{}"
# Mapping from the node type tag to the section of the autoscaler yaml that
# contains the config for the node type.
NODE_KIND_CONFIG_KEYS = {
NODE_KIND_WORKER: "worker_nodes",
NODE_KIND_HEAD: "head_node",
}
DEFAULT_AMI_NAME = "AWS Deep Learning AMI (Ubuntu 18.04) V30.0"
# Obtained from https://aws.amazon.com/marketplace/pp/B07Y43P7X5 on 8/4/2020.
DEFAULT_AMI = {
"us-east-1": "ami-029510cec6d69f121", # US East (N. Virginia)
"us-east-2": "ami-08bf49c7b3a0c761e", # US East (Ohio)
"us-west-1": "ami-0cc472544ce594a19", # US West (N. California)
"us-west-2": "ami-0a2363a9cff180a64", # US West (Oregon)
"ca-central-1": "ami-0a871851b2ab39f01", # Canada (Central)
"eu-central-1": "ami-049fb1ea198d189d7", # EU (Frankfurt)
"eu-west-1": "ami-0abcbc65f89fb220e", # EU (Ireland)
"eu-west-2": "ami-0755b39fd4dab7cbe", # EU (London)
"eu-west-3": "ami-020485d8df1d45530", # EU (Paris)
"sa-east-1": "ami-058a6883cbdb4e599", # SA (Sao Paulo)
}
# todo: cli_logger should handle this assert properly
# this should probably also happens somewhere else
assert StrictVersion(boto3.__version__) >= StrictVersion("1.4.8"), \
"Boto3 version >= 1.4.8 required, try `pip install -U boto3`"
def key_pair(i, region, key_name):
"""
If key_name is not None, key_pair will be named after key_name.
Returns the ith default (aws_key_pair_name, key_pair_path).
"""
if i == 0:
key_pair_name = ("{}_{}".format(RAY, region)
if key_name is None else key_name)
return (key_pair_name,
os.path.expanduser("~/.ssh/{}.pem".format(key_pair_name)))
key_pair_name = ("{}_{}_{}".format(RAY, i, region)
if key_name is None else key_name + "_key-{}".format(i))
return (key_pair_name,
os.path.expanduser("~/.ssh/{}.pem".format(key_pair_name)))
# Suppress excessive connection dropped logs from boto
logging.getLogger("botocore").setLevel(logging.WARNING)
_log_info = {}
def reload_log_state(override_log_info):
_log_info.update(override_log_info)
def get_log_state():
return _log_info.copy()
def _set_config_info(**kwargs):
"""Record configuration artifacts useful for logging."""
# todo: this is technically fragile iff we ever use multiple configs
for k, v in kwargs.items():
_log_info[k] = v
def _arn_to_name(arn):
return arn.split(":")[-1].split("/")[-1]
def log_to_cli(config):
provider_name = PROVIDER_PRETTY_NAMES.get("aws", None)
cli_logger.doassert(provider_name is not None,
"Could not find a pretty name for the AWS provider.")
with cli_logger.group("{} config", provider_name):
def same_everywhere(key):
return config["head_node"][key] == config["worker_nodes"][key]
def print_info(resource_string,
key,
head_src_key,
workers_src_key,
allowed_tags=["default"],
list_value=False):
head_tags = {}
workers_tags = {}
if _log_info[head_src_key] in allowed_tags:
head_tags[_log_info[head_src_key]] = True
if _log_info[workers_src_key] in allowed_tags:
workers_tags[_log_info[workers_src_key]] = True
head_value_str = config["head_node"][key]
if list_value:
head_value_str = cli_logger.render_list(head_value_str)
if same_everywhere(key):
cli_logger.labeled_value( # todo: handle plural vs singular?
resource_string + " (head & workers)",
"{}",
head_value_str,
_tags=head_tags)
else:
workers_value_str = config["worker_nodes"][key]
if list_value:
workers_value_str = cli_logger.render_list(
workers_value_str)
cli_logger.labeled_value(
resource_string + " (head)",
"{}",
head_value_str,
_tags=head_tags)
cli_logger.labeled_value(
resource_string + " (workers)",
"{}",
workers_value_str,
_tags=workers_tags)
tags = {"default": _log_info["head_instance_profile_src"] == "default"}
cli_logger.labeled_value(
"IAM Profile",
"{}",
_arn_to_name(config["head_node"]["IamInstanceProfile"]["Arn"]),
_tags=tags)
if ("KeyName" in config["head_node"]
and "KeyName" in config["worker_nodes"]):
print_info("EC2 Key pair", "KeyName", "keypair_src", "keypair_src")
print_info(
"VPC Subnets",
"SubnetIds",
"head_subnet_src",
"workers_subnet_src",
list_value=True)
print_info(
"EC2 Security groups",
"SecurityGroupIds",
"head_security_group_src",
"workers_security_group_src",
list_value=True)
print_info(
"EC2 AMI",
"ImageId",
"head_ami_src",
"workers_ami_src",
allowed_tags=["dlami"])
cli_logger.newline()
def bootstrap_aws(config):
# The head node needs to have an IAM role that allows it to create further
# EC2 instances.
config = _configure_iam_role(config)
# Configure SSH access, using an existing key pair if possible.
config = _configure_key_pair(config)
# Pick a reasonable subnet if not specified by the user.
config = _configure_subnet(config)
# Cluster workers should be in a security group that permits traffic within
# the group, and also SSH access from outside.
config = _configure_security_group(config)
# Provide a helpful message for missing AMI.
_check_ami(config)
return config
def _configure_iam_role(config):
if "IamInstanceProfile" in config["head_node"]:
_set_config_info(head_instance_profile_src="config")
return config
_set_config_info(head_instance_profile_src="default")
profile = _get_instance_profile(DEFAULT_RAY_INSTANCE_PROFILE, config)
if profile is None:
cli_logger.verbose(
"Creating new IAM instance profile {} for use as the default.",
cf.bold(DEFAULT_RAY_INSTANCE_PROFILE))
cli_logger.old_info(
logger, "_configure_iam_role: "
"Creating new instance profile {}", DEFAULT_RAY_INSTANCE_PROFILE)
client = _client("iam", config)
client.create_instance_profile(
InstanceProfileName=DEFAULT_RAY_INSTANCE_PROFILE)
profile = _get_instance_profile(DEFAULT_RAY_INSTANCE_PROFILE, config)
time.sleep(15) # wait for propagation
cli_logger.doassert(profile is not None,
"Failed to create instance profile.") # todo: err msg
assert profile is not None, "Failed to create instance profile"
if not profile.roles:
role = _get_role(DEFAULT_RAY_IAM_ROLE, config)
if role is None:
cli_logger.verbose(
"Creating new IAM role {} for "
"use as the default instance role.",
cf.bold(DEFAULT_RAY_IAM_ROLE))
cli_logger.old_info(logger, "_configure_iam_role: "
"Creating new role {}", DEFAULT_RAY_IAM_ROLE)
iam = _resource("iam", config)
iam.create_role(
RoleName=DEFAULT_RAY_IAM_ROLE,
AssumeRolePolicyDocument=json.dumps({
"Statement": [
{
"Effect": "Allow",
"Principal": {
"Service": "ec2.amazonaws.com"
},
"Action": "sts:AssumeRole",
},
],
}))
role = _get_role(DEFAULT_RAY_IAM_ROLE, config)
cli_logger.doassert(role is not None,
"Failed to create role.") # todo: err msg
assert role is not None, "Failed to create role"
role.attach_policy(
PolicyArn="arn:aws:iam::aws:policy/AmazonEC2FullAccess")
role.attach_policy(
PolicyArn="arn:aws:iam::aws:policy/AmazonS3FullAccess")
profile.add_role(RoleName=role.name)
time.sleep(15) # wait for propagation
cli_logger.old_info(
logger, "_configure_iam_role: "
"Role not specified for head node, using {}", profile.arn)
config["head_node"]["IamInstanceProfile"] = {"Arn": profile.arn}
return config
def _configure_key_pair(config):
if "ssh_private_key" in config["auth"]:
_set_config_info(keypair_src="config")
# If the key is not configured via the cloudinit
# UserData, it should be configured via KeyName or
# else we will risk starting a node that we cannot
# SSH into:
if "UserData" not in config["head_node"]:
cli_logger.doassert( # todo: verify schema beforehand?
"KeyName" in config["head_node"],
"`KeyName` missing for head node.") # todo: err msg
assert "KeyName" in config["head_node"]
if "UserData" not in config["worker_nodes"]:
cli_logger.doassert(
"KeyName" in config["worker_nodes"],
"`KeyName` missing for worker nodes.") # todo: err msg
assert "KeyName" in config["worker_nodes"]
return config
_set_config_info(keypair_src="default")
ec2 = _resource("ec2", config)
# Writing the new ssh key to the filesystem fails if the ~/.ssh
# directory doesn't already exist.
os.makedirs(os.path.expanduser("~/.ssh"), exist_ok=True)
# Try a few times to get or create a good key pair.
MAX_NUM_KEYS = 30
for i in range(MAX_NUM_KEYS):
key_name = config["provider"].get("key_pair", {}).get("key_name")
key_name, key_path = key_pair(i, config["provider"]["region"],
key_name)
key = _get_key(key_name, config)
# Found a good key.
if key and os.path.exists(key_path):
break
# We can safely create a new key.
if not key and not os.path.exists(key_path):
cli_logger.verbose(
"Creating new key pair {} for use as the default.",
cf.bold(key_name))
cli_logger.old_info(
logger, "_configure_key_pair: "
"Creating new key pair {}", key_name)
key = ec2.create_key_pair(KeyName=key_name)
# We need to make sure to _create_ the file with the right
# permissions. In order to do that we need to change the default
# os.open behavior to include the mode we want.
with open(key_path, "w", opener=partial(os.open, mode=0o600)) as f:
f.write(key.key_material)
break
if not key:
cli_logger.abort(
"No matching local key file for any of the key pairs in this "
"account with ids from 0..{}. "
"Consider deleting some unused keys pairs from your account.",
key_name) # todo: err msg
raise ValueError(
"No matching local key file for any of the key pairs in this "
"account with ids from 0..{}. ".format(key_name) +
"Consider deleting some unused keys pairs from your account.")
cli_logger.doassert(
os.path.exists(key_path), "Private key file " + cf.bold("{}") +
" not found for " + cf.bold("{}"), key_path, key_name) # todo: err msg
assert os.path.exists(key_path), \
"Private key file {} not found for {}".format(key_path, key_name)
cli_logger.old_info(
logger, "_configure_key_pair: "
"KeyName not specified for nodes, using {}", key_name)
config["auth"]["ssh_private_key"] = key_path
config["head_node"]["KeyName"] = key_name
config["worker_nodes"]["KeyName"] = key_name
return config
def _configure_subnet(config):
ec2 = _resource("ec2", config)
use_internal_ips = config["provider"].get("use_internal_ips", False)
try:
subnets = sorted(
(s for s in ec2.subnets.all() if s.state == "available" and (
use_internal_ips or s.map_public_ip_on_launch)),
reverse=True, # sort from Z-A
key=lambda subnet: subnet.availability_zone)
except botocore.exceptions.ClientError as exc:
handle_boto_error(exc, "Failed to fetch available subnets from AWS.")
raise exc
if not subnets:
cli_logger.abort(
"No usable subnets found, try manually creating an instance in "
"your specified region to populate the list of subnets "
"and trying this again.\n"
"Note that the subnet must map public IPs "
"on instance launch unless you set `use_internal_ips: true` in "
"the `provider` config.") # todo: err msg
raise Exception(
"No usable subnets found, try manually creating an instance in "
"your specified region to populate the list of subnets "
"and trying this again. Note that the subnet must map public IPs "
"on instance launch unless you set 'use_internal_ips': True in "
"the 'provider' config.")
if "availability_zone" in config["provider"]:
azs = config["provider"]["availability_zone"].split(",")
subnets = [s for s in subnets if s.availability_zone in azs]
if not subnets:
cli_logger.abort(
"No usable subnets matching availability zone {} found.\n"
"Choose a different availability zone or try "
"manually creating an instance in your specified region "
"to populate the list of subnets and trying this again.",
config["provider"]["availability_zone"]) # todo: err msg
raise Exception(
"No usable subnets matching availability zone {} "
"found. Choose a different availability zone or try "
"manually creating an instance in your specified region "
"to populate the list of subnets and trying this again.".
format(config["provider"]["availability_zone"]))
subnet_ids = [s.subnet_id for s in subnets]
subnet_descr = [(s.subnet_id, s.availability_zone) for s in subnets]
if "SubnetIds" not in config["head_node"]:
_set_config_info(head_subnet_src="default")
config["head_node"]["SubnetIds"] = subnet_ids
cli_logger.old_info(
logger, "_configure_subnet: "
"SubnetIds not specified for head node, using {}", subnet_descr)
else:
_set_config_info(head_subnet_src="config")
if "SubnetIds" not in config["worker_nodes"]:
_set_config_info(workers_subnet_src="default")
config["worker_nodes"]["SubnetIds"] = subnet_ids
cli_logger.old_info(
logger, "_configure_subnet: "
"SubnetId not specified for workers,"
" using {}", subnet_descr)
else:
_set_config_info(workers_subnet_src="config")
return config
def _configure_security_group(config):
_set_config_info(
head_security_group_src="config", workers_security_group_src="config")
node_types_to_configure = [
node_type for node_type, config_key in NODE_KIND_CONFIG_KEYS.items()
if "SecurityGroupIds" not in config[NODE_KIND_CONFIG_KEYS[node_type]]
]
if not node_types_to_configure:
return config # have user-defined groups
security_groups = _upsert_security_groups(config, node_types_to_configure)
if NODE_KIND_HEAD in node_types_to_configure:
head_sg = security_groups[NODE_KIND_HEAD]
_set_config_info(head_security_group_src="default")
cli_logger.old_info(
logger, "_configure_security_group: "
"SecurityGroupIds not specified for head node, using {} ({})",
head_sg.group_name, head_sg.id)
config["head_node"]["SecurityGroupIds"] = [head_sg.id]
if NODE_KIND_WORKER in node_types_to_configure:
workers_sg = security_groups[NODE_KIND_WORKER]
_set_config_info(workers_security_group_src="default")
cli_logger.old_info(
logger, "_configure_security_group: "
"SecurityGroupIds not specified for workers, using {} ({})",
workers_sg.group_name, workers_sg.id)
config["worker_nodes"]["SecurityGroupIds"] = [workers_sg.id]
return config
def _check_ami(config):
"""Provide helpful message for missing ImageId for node configuration."""
_set_config_info(head_ami_src="config", workers_ami_src="config")
region = config["provider"]["region"]
default_ami = DEFAULT_AMI.get(region)
if not default_ami:
# If we do not provide a default AMI for the given region, noop.
return
if config["head_node"].get("ImageId", "").lower() == "latest_dlami":
config["head_node"]["ImageId"] = default_ami
_set_config_info(head_ami_src="dlami")
cli_logger.old_info(
logger,
"_check_ami: head node ImageId is 'latest_dlami'. "
"Using '{ami_id}', which is the default {ami_name} "
"for your region ({region}).",
ami_id=default_ami,
ami_name=DEFAULT_AMI_NAME,
region=region)
if config["worker_nodes"].get("ImageId", "").lower() == "latest_dlami":
config["worker_nodes"]["ImageId"] = default_ami
_set_config_info(workers_ami_src="dlami")
cli_logger.old_info(
logger,
"_check_ami: worker nodes ImageId is 'latest_dlami'. "
"Using '{ami_id}', which is the default {ami_name} "
"for your region ({region}).",
ami_id=default_ami,
ami_name=DEFAULT_AMI_NAME,
region=region)
def _upsert_security_groups(config, node_types):
security_groups = _get_or_create_vpc_security_groups(config, node_types)
_upsert_security_group_rules(config, security_groups)
return security_groups
def _get_or_create_vpc_security_groups(conf, node_types):
# Figure out which VPC each node_type is in...
ec2 = _resource("ec2", conf)
node_type_to_vpc = {
node_type: _get_vpc_id_or_die(
ec2,
conf[NODE_KIND_CONFIG_KEYS[node_type]]["SubnetIds"][0],
)
for node_type in node_types
}
# Generate the name of the security group we're looking for...
expected_sg_name = SECURITY_GROUP_TEMPLATE.format(conf["cluster_name"])
# Figure out which security groups with this name exist for each VPC...
vpc_to_existing_sg = {
sg.vpc_id: sg
for sg in _get_security_groups(
conf,
node_type_to_vpc.values(),
[expected_sg_name],
)
}
# Lazily create any security group we're missing for each VPC...
vpc_to_sg = LazyDefaultDict(
partial(_create_security_group, conf, group_name=expected_sg_name),
vpc_to_existing_sg,
)
# Then return a mapping from each node_type to its security group...
return {
node_type: vpc_to_sg[vpc_id]
for node_type, vpc_id in node_type_to_vpc.items()
}
@lru_cache()
def _get_vpc_id_or_die(ec2, subnet_id):
subnet = list(
ec2.subnets.filter(Filters=[{
"Name": "subnet-id",
"Values": [subnet_id]
}]))
# TODO: better error message
cli_logger.doassert(len(subnet) == 1, "Subnet ID not found: {}", subnet_id)
assert len(subnet) == 1, "Subnet ID not found: {}".format(subnet_id)
subnet = subnet[0]
return subnet.vpc_id
def _get_security_group(config, vpc_id, group_name):
security_group = _get_security_groups(config, [vpc_id], [group_name])
return None if not security_group else security_group[0]
def _get_security_groups(config, vpc_ids, group_names):
unique_vpc_ids = list(set(vpc_ids))
unique_group_names = set(group_names)
ec2 = _resource("ec2", config)
existing_groups = list(
ec2.security_groups.filter(Filters=[{
"Name": "vpc-id",
"Values": unique_vpc_ids
}]))
filtered_groups = [
sg for sg in existing_groups if sg.group_name in unique_group_names
]
return filtered_groups
def _create_security_group(config, vpc_id, group_name):
client = _client("ec2", config)
client.create_security_group(
Description="Auto-created security group for Ray workers",
GroupName=group_name,
VpcId=vpc_id)
security_group = _get_security_group(config, vpc_id, group_name)
cli_logger.verbose(
"Created new security group {}",
cf.bold(security_group.group_name),
_tags=dict(id=security_group.id))
cli_logger.old_info(
logger, "_create_security_group: Created new security group {} ({})",
security_group.group_name, security_group.id)
cli_logger.doassert(security_group,
"Failed to create security group") # err msg
assert security_group, "Failed to create security group"
return security_group
def _upsert_security_group_rules(conf, security_groups):
sgids = {sg.id for sg in security_groups.values()}
# sort security group items for deterministic inbound rule config order
# (mainly supports more precise stub-based boto3 unit testing)
for node_type, sg in sorted(security_groups.items()):
sg = security_groups[node_type]
if not sg.ip_permissions:
_update_inbound_rules(sg, sgids, conf)
def _update_inbound_rules(target_security_group, sgids, config):
extended_rules = config["provider"] \
.get("security_group", {}) \
.get("IpPermissions", [])
ip_permissions = _create_default_inbound_rules(sgids, extended_rules)
target_security_group.authorize_ingress(IpPermissions=ip_permissions)
def _create_default_inbound_rules(sgids, extended_rules=[]):
intracluster_rules = _create_default_instracluster_inbound_rules(sgids)
ssh_rules = _create_default_ssh_inbound_rules()
merged_rules = itertools.chain(
intracluster_rules,
ssh_rules,
extended_rules,
)
return list(merged_rules)
def _create_default_instracluster_inbound_rules(intracluster_sgids):
return [{
"FromPort": -1,
"ToPort": -1,
"IpProtocol": "-1",
"UserIdGroupPairs": [
{
"GroupId": security_group_id
} for security_group_id in sorted(intracluster_sgids)
# sort security group IDs for deterministic IpPermission models
# (mainly supports more precise stub-based boto3 unit testing)
]
}]
def _create_default_ssh_inbound_rules():
return [{
"FromPort": 22,
"ToPort": 22,
"IpProtocol": "tcp",
"IpRanges": [{
"CidrIp": "0.0.0.0/0"
}]
}]
def _get_role(role_name, config):
iam = _resource("iam", config)
role = iam.Role(role_name)
try:
role.load()
return role
except botocore.exceptions.ClientError as exc:
if exc.response.get("Error", {}).get("Code") == "NoSuchEntity":
return None
else:
handle_boto_error(
exc, "Failed to fetch IAM role data for {} from AWS.",
cf.bold(role_name))
raise exc
def _get_instance_profile(profile_name, config):
iam = _resource("iam", config)
profile = iam.InstanceProfile(profile_name)
try:
profile.load()
return profile
except botocore.exceptions.ClientError as exc:
if exc.response.get("Error", {}).get("Code") == "NoSuchEntity":
return None
else:
handle_boto_error(
exc,
"Failed to fetch IAM instance profile data for {} from AWS.",
cf.bold(profile_name))
raise exc
def _get_key(key_name, config):
ec2 = _resource("ec2", config)
try:
for key in ec2.key_pairs.filter(Filters=[{
"Name": "key-name",
"Values": [key_name]
}]):
if key.name == key_name:
return key
except botocore.exceptions.ClientError as exc:
handle_boto_error(exc, "Failed to fetch EC2 key pair {} from AWS.",
cf.bold(key_name))
raise exc
def _client(name, config):
return _resource(name, config).meta.client
def _resource(name, config):
region = config["provider"]["region"]
aws_credentials = config["provider"].get("aws_credentials", {})
return _resource_cache(name, region, **aws_credentials)
@lru_cache()
def _resource_cache(name, region, **kwargs):
boto_config = Config(retries={"max_attempts": BOTO_MAX_RETRIES})
return boto3.resource(
name,
region,
config=boto_config,
**kwargs,
)
python/ray/autoscaler/aws/node_provider.py
-481
View File
@@ -1,481 +0,0 @@
import random
import copy
import threading
from collections import defaultdict
import logging
import boto3
import botocore
from botocore.config import Config
from ray.autoscaler.node_provider import NodeProvider
from ray.autoscaler.aws.config import bootstrap_aws
from ray.autoscaler.tags import TAG_RAY_CLUSTER_NAME, TAG_RAY_NODE_NAME, \
TAG_RAY_LAUNCH_CONFIG, TAG_RAY_NODE_KIND, TAG_RAY_USER_NODE_TYPE
from ray.ray_constants import BOTO_MAX_RETRIES, BOTO_CREATE_MAX_RETRIES
from ray.autoscaler.log_timer import LogTimer
from ray.autoscaler.aws.utils import boto_exception_handler
from ray.autoscaler.cli_logger import cli_logger
import colorful as cf
logger = logging.getLogger(__name__)
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_client(region, max_retries, aws_credentials=None):
"""Make client, retrying requests up to `max_retries`."""
config = Config(retries={"max_attempts": max_retries})
aws_credentials = aws_credentials or {}
return boto3.resource(
"ec2", region_name=region, config=config, **aws_credentials)
class AWSNodeProvider(NodeProvider):
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_client(
region=provider_config["region"],
max_retries=BOTO_MAX_RETRIES,
aws_credentials=aws_credentials)
self.ec2_fail_fast = make_ec2_client(
region=provider_config["region"],
max_retries=0,
aws_credentials=aws_credentials)
# Try availability zones round-robin, starting from random offset
self.subnet_idx = random.randint(0, 100)
self.tag_cache = {} # Tags that we believe to actually be on EC2.
self.tag_cache_pending = {} # Tags that we will soon upload.
self.tag_cache_lock = threading.Lock()
self.tag_cache_update_event = threading.Event()
self.tag_cache_kill_event = threading.Event()
self.tag_update_thread = threading.Thread(
target=self._node_tag_update_loop)
self.tag_update_thread.start()
# Cache of node objects from the last nodes() call. This avoids
# excessive DescribeInstances requests.
self.cached_nodes = {}
def _node_tag_update_loop(self):
"""Update the AWS tags for a cluster periodically.
The purpose of this loop is to avoid excessive EC2 calls when a large
number of nodes are being launched simultaneously.
"""
while True:
self.tag_cache_update_event.wait()
self.tag_cache_update_event.clear()
batch_updates = defaultdict(list)
with self.tag_cache_lock:
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 = {}
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
}],
)
self.tag_cache_kill_event.wait(timeout=5)
if self.tag_cache_kill_event.is_set():
return
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):
with self.tag_cache_lock:
try:
self.tag_cache_pending[node_id].update(tags)
except KeyError:
self.tag_cache_pending[node_id] = tags
self.tag_cache_update_event.set()
def create_node(self, node_config, tags, count):
tags = copy.deepcopy(tags)
# 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]],
})
reuse_nodes = list(
self.ec2.instances.filter(Filters=filters))[:count]
reuse_node_ids = [n.id 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))
cli_logger.old_info(
logger, "AWSNodeProvider: reusing instances {}. "
"To disable reuse, set "
"'cache_stopped_nodes: False' in the provider "
"config.", 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)
cli_logger.old_info(
logger,
"AWSNodeProvider: waiting for instance "
"{} to fully 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)
if count:
self._create_node(node_config, tags, count)
def _create_node(self, node_config, tags, count):
tags = to_aws_format(tags)
conf = node_config.copy()
# Delete unsupported keys from the node config
try:
del conf["Resources"]
except KeyError:
pass
tag_pairs = [{
"Key": TAG_RAY_CLUSTER_NAME,
"Value": self.cluster_name,
}]
for k, v in tags.items():
tag_pairs.append({
"Key": k,
"Value": v,
})
tag_specs = [{
"ResourceType": "instance",
"Tags": tag_pairs,
}]
user_tag_specs = conf.get("TagSpecifications", [])
# Allow users to add tags and override values of existing
# tags with their own. This only applies to the resource type
# "instance". All other resource types are appended to the list of
# 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]
# SubnetIds is not a real config key: we must resolve to a
# single SubnetId before invoking the AWS API.
subnet_ids = conf.pop("SubnetIds")
for attempt in range(1, BOTO_CREATE_MAX_RETRIES + 1):
try:
subnet_id = subnet_ids[self.subnet_idx % len(subnet_ids)]
cli_logger.old_info(
logger, "NodeProvider: calling create_instances "
"with {} (count={}).", subnet_id, count)
self.subnet_idx += 1
conf.update({
"MinCount": 1,
"MaxCount": count,
"SubnetId": subnet_id,
"TagSpecifications": tag_specs
})
created = self.ec2_fail_fast.create_instances(**conf)
# todo: timed?
# todo: handle plurality?
with cli_logger.group(
"Launched {} nodes",
count,
_tags=dict(subnet_id=subnet_id)):
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 = instance.state_reason or {
"Message": "pending"
}
cli_logger.print(
"Launched instance {}",
instance.instance_id,
_tags=dict(
state=instance.state["Name"],
info=state_reason["Message"]))
cli_logger.old_info(
logger, "NodeProvider: Created instance "
"[id={}, name={}, info={}]", instance.instance_id,
instance.state["Name"], state_reason["Message"])
break
except botocore.exceptions.ClientError as exc:
if attempt == BOTO_CREATE_MAX_RETRIES:
# todo: err msg
cli_logger.abort(
"Failed to launch instances. Max attempts exceeded.")
cli_logger.old_error(
logger,
"create_instances: Max attempts ({}) exceeded.",
BOTO_CREATE_MAX_RETRIES)
raise exc
else:
cli_logger.print(
"create_instances: Attempt failed with {}, retrying.",
exc)
cli_logger.old_error(logger, exc)
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?
cli_logger.old_info(
logger,
"AWSNodeProvider: terminating node {} (spot nodes cannot "
"be stopped, only terminated)", node_id)
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?
cli_logger.old_info(
logger,
"AWSNodeProvider: stopping node {}. To terminate nodes "
"on stop, set 'cache_stopped_nodes: False' in the "
"provider config.".format(node_id))
node.stop()
else:
node.terminate()
self.tag_cache.pop(node_id, None)
self.tag_cache_pending.pop(node_id, None)
def terminate_nodes(self, node_ids):
if not node_ids:
return
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))
cli_logger.old_info(
logger,
"AWSNodeProvider: stopping nodes {}. To terminate nodes "
"on stop, set 'cache_stopped_nodes: False' in the "
"provider config.", on_demand_ids)
self.ec2.meta.client.stop_instances(InstanceIds=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))
cli_logger.old_info(
logger,
"AWSNodeProvider: terminating nodes {} (spot nodes cannot "
"be stopped, only terminated)", spot_ids)
self.ec2.meta.client.terminate_instances(InstanceIds=spot_ids)
else:
self.ec2.meta.client.terminate_instances(InstanceIds=node_ids)
for node_id in node_ids:
self.tag_cache.pop(node_id, None)
self.tag_cache_pending.pop(node_id, None)
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.
matches = list(self.ec2.instances.filter(InstanceIds=[node_id]))
assert len(matches) == 1, "Invalid instance id {}".format(node_id)
return matches[0]
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)
def cleanup(self):
self.tag_cache_update_event.set()
self.tag_cache_kill_event.set()
@staticmethod
def bootstrap_config(cluster_config):
return bootstrap_aws(cluster_config)
python/ray/autoscaler/aws/utils.py
-128
View File
@@ -1,128 +0,0 @@
from collections import defaultdict
from ray.autoscaler.cli_logger import cli_logger
import colorful as cf
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):
if cli_logger.old_style:
# old-style logging doesn't do anything here
# so we exit early
return
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 = [
"{}\n"
"Error 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()