def test_fills_out_amis(iam_client_stub, ec2_client_stub):
# Setup stubs to mock out boto3
stubs.configure_iam_role_default(iam_client_stub)
stubs.configure_key_pair_default(ec2_client_stub)
stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG)
stubs.configure_subnet_default(ec2_client_stub)
config = helpers.load_aws_example_config_file("example-full.yaml")
del config["available_node_types"]["ray.head.default"]["node_config"][
"ImageId"]
del config["available_node_types"]["ray.worker.default"]["node_config"][
"ImageId"]
# Pass in SG for stub to work
config["head_node"]["SecurityGroupIds"] = ["sg-1234abcd"]
config["worker_nodes"]["SecurityGroupIds"] = ["sg-1234abcd"]
defaults_filled = bootstrap_aws(config)
ami = DEFAULT_AMI.get(config.get("provider", {}).get("region"))
assert defaults_filled["head_node"].get("ImageId") == ami
assert defaults_filled["worker_nodes"].get("ImageId") == ami
iam_client_stub.assert_no_pending_responses()
ec2_client_stub.assert_no_pending_responses()
def test_iam_already_configured(iam_client_stub, ec2_client_stub):
"""
Checks that things work as expected when IAM role is supplied by user.
"""
stubs.configure_key_pair_default(ec2_client_stub)
stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG)
stubs.configure_subnet_default(ec2_client_stub)
config = helpers.load_aws_example_config_file("example-full.yaml")
head_node_config = config["available_node_types"]["ray.head.default"][
"node_config"]
worker_node_config = config["available_node_types"]["ray.worker.default"][
"node_config"]
head_node_config["IamInstanceProfile"] = "mock_profile"
# Pass in SG for stub to work
head_node_config["SecurityGroupIds"] = ["sg-1234abcd"]
worker_node_config["SecurityGroupIds"] = ["sg-1234abcd"]
defaults_filled = bootstrap_aws(config)
filled_head = defaults_filled["available_node_types"]["ray.head.default"][
"node_config"]
assert filled_head["IamInstanceProfile"] == "mock_profile"
assert "IamInstanceProfile" not in defaults_filled["head_node"]
iam_client_stub.assert_no_pending_responses()
ec2_client_stub.assert_no_pending_responses()
def test_network_interface_missing_security_group():
# If NetworkInterfaces are defined, each must have security groups
expected_error_msg = ("NetworkInterfaces are defined but at least one is "
"missing a security group. Please ensure all "
"interfaces have a security group assigned.")
config = helpers.load_aws_example_config_file(
"example-network-interfaces.yaml")
for name, node_type in config["available_node_types"].items():
node_cfg = node_type["node_config"]
for network_interface_cfg in node_cfg["NetworkInterfaces"]:
network_interface_cfg.pop("Groups")
with pytest.raises(ValueError, match=expected_error_msg):
helpers.bootstrap_aws_config(config)
def test_network_interface_missing_subnet():
# If NetworkInterfaces are defined, each must have a subnet ID
expected_error_msg = "NetworkInterfaces are defined but at least one is " \
"missing a subnet. Please ensure all interfaces " \
"have a subnet assigned."
config = helpers.load_aws_example_config_file(
"example-network-interfaces.yaml")
for name, node_type in config["available_node_types"].items():
node_cfg = node_type["node_config"]
for network_interface_cfg in node_cfg["NetworkInterfaces"]:
network_interface_cfg.pop("SubnetId")
with pytest.raises(ValueError, match=expected_error_msg):
helpers.bootstrap_aws_config(config)
def test_fills_out_amis_and_iam(iam_client_stub, ec2_client_stub, region):
# Set up expected key pair for specific region
region_key_pair = DEFAULT_KEY_PAIR.copy()
region_key_pair["KeyName"] = DEFAULT_KEY_PAIR["KeyName"].replace(
"us-west-2", region)
# Setup stubs to mock out boto3
stubs.configure_iam_role_default(iam_client_stub)
stubs.configure_key_pair_default(ec2_client_stub,
region=region,
expected_key_pair=region_key_pair)
stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG)
stubs.configure_subnet_default(ec2_client_stub)
config = helpers.load_aws_example_config_file("example-full.yaml")
head_node_config = config["available_node_types"]["ray.head.default"][
"node_config"]
worker_node_config = config["available_node_types"]["ray.worker.default"][
"node_config"]
del head_node_config["ImageId"]
del worker_node_config["ImageId"]
# Pass in SG for stub to work
head_node_config["SecurityGroupIds"] = ["sg-1234abcd"]
worker_node_config["SecurityGroupIds"] = ["sg-1234abcd"]
config["provider"]["region"] = region
defaults_filled = bootstrap_aws(config)
ami = DEFAULT_AMI.get(defaults_filled.get("provider", {}).get("region"))
for node_type in defaults_filled["available_node_types"].values():
node_config = node_type["node_config"]
assert node_config.get("ImageId") == ami
# Correctly configured IAM role
assert defaults_filled["head_node"]["IamInstanceProfile"] == {
"Arn": DEFAULT_INSTANCE_PROFILE["Arn"]
}
# Workers of the head's type do not get the IAM role.
head_type = config["head_node_type"]
assert ("IamInstanceProfile"
not in defaults_filled["available_node_types"][head_type])
iam_client_stub.assert_no_pending_responses()
ec2_client_stub.assert_no_pending_responses()
def test_network_interface_conflict_keys():
# If NetworkInterfaces are defined, SubnetId and SecurityGroupIds
# can't be specified in the same node type config.
conflict_kv_pairs = [("SubnetId", "subnet-0000000"),
("SubnetIds", ["subnet-0000000", "subnet-1111111"]),
("SecurityGroupIds", ["sg-1234abcd", "sg-dcba4321"])]
expected_error_msg = "If NetworkInterfaces are defined, subnets and " \
"security groups must ONLY be given in each " \
"NetworkInterface."
for conflict_kv_pair in conflict_kv_pairs:
config = helpers.load_aws_example_config_file(
"example-network-interfaces.yaml")
head_name = config["head_node_type"]
head_node_cfg = config["available_node_types"][head_name][
"node_config"]
head_node_cfg[conflict_kv_pair[0]] = conflict_kv_pair[1]
with pytest.raises(ValueError, match=expected_error_msg):
helpers.bootstrap_aws_config(config)
def test_use_subnets_ordered_by_az(ec2_client_stub):
"""
This test validates that when bootstrap_aws populates the SubnetIds field,
the subnets are ordered the same way as availability zones.
"""
# Add a response with a twenty subnets round-robined across the 4 AZs in
# `us-west-2` (a,b,c,d). At the end we should only have 15 subnets, ordered
# first from `us-west-2c`, then `us-west-2d`, then `us-west-2a`.
stubs.describe_twenty_subnets_in_different_azs(ec2_client_stub)
base_config = helpers.load_aws_example_config_file("example-full.yaml")
base_config["provider"][
"availability_zone"] = "us-west-2c,us-west-2d,us-west-2a"
config = _configure_subnet(base_config)
# We've filtered down to only subnets in 2c, 2d & 2a
for node_type in config["available_node_types"].values():
node_config = node_type["node_config"]
assert len(node_config["SubnetIds"]) == 15
offsets = [int(s.split("-")[1]) % 4 for s in node_config["SubnetIds"]]
assert set(offsets[:5]) == {2}, "First 5 should be in us-west-2c"
assert set(offsets[5:10]) == {3}, "Next 5 should be in us-west-2d"
assert set(offsets[10:15]) == {0}, "Last 5 should be in us-west-2a"
def test_missing_keyname(iam_client_stub, ec2_client_stub):
config = helpers.load_aws_example_config_file("example-full.yaml")
config["auth"]["ssh_private_key"] = "/path/to/private/key"
head_node_config = config["available_node_types"]["ray.head.default"][
"node_config"]
worker_node_config = config["available_node_types"]["ray.worker.default"][
"node_config"]
# Setup stubs to mock out boto3. Should fail on assertion after
# checking KeyName/UserData.
stubs.configure_iam_role_default(iam_client_stub)
missing_user_data_config = copy.deepcopy(config)
with pytest.raises(AssertionError):
# Config specified ssh_private_key, but missing KeyName/UserData in
# node configs
bootstrap_aws(missing_user_data_config)
# Pass in SG for stub to work
head_node_config["SecurityGroupIds"] = ["sg-1234abcd"]
worker_node_config["SecurityGroupIds"] = ["sg-1234abcd"]
# Set UserData for both node configs
head_node_config["UserData"] = {"someKey": "someValue"}
worker_node_config["UserData"] = {"someKey": "someValue"}
# Stubs to mock out boto3. Should no longer fail on assertion
# and go on to describe security groups + configure subnet
stubs.configure_iam_role_default(iam_client_stub)
stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG)
stubs.configure_subnet_default(ec2_client_stub)
# Should work without error now that UserData is set
bootstrap_aws(config)
iam_client_stub.assert_no_pending_responses()
ec2_client_stub.assert_no_pending_responses()
def test_log_to_cli(iam_client_stub, ec2_client_stub):
config = helpers.load_aws_example_config_file("example-full.yaml")
head_node_config = config["available_node_types"]["ray.head.default"][
"node_config"]
worker_node_config = config["available_node_types"]["ray.worker.default"][
"node_config"]
# Pass in SG for stub to work
head_node_config["SecurityGroupIds"] = ["sg-1234abcd"]
worker_node_config["SecurityGroupIds"] = ["sg-1234abcd"]
stubs.configure_iam_role_default(iam_client_stub)
stubs.configure_key_pair_default(ec2_client_stub)
stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG)
stubs.configure_subnet_default(ec2_client_stub)
config = helpers.bootstrap_aws_config(config)
# Only side-effect is to print logs to cli, called just to
# check that it runs without error
log_to_cli(config)
iam_client_stub.assert_no_pending_responses()
ec2_client_stub.assert_no_pending_responses()
def test_create_sg_multinode(iam_client_stub, ec2_client_stub):
"""
Test AWS Bootstrap logic when config being bootstrapped has the
following properties:
(1) auth config does not specify ssh key path
(2) available_node_types is provided
(3) security group name and ip permissions set in provider field
(4) Available node types have SubnetIds field set and this
field is of form SubnetIds: [subnet-xxxxx].
Both node types specify the same subnet-xxxxx.
Tests creation of a security group and key pair under these conditions.
"""
# Generate a config of the desired form.
subnet_id = DEFAULT_SUBNET["SubnetId"]
# security group info to go in provider field
provider_data = helpers.load_aws_example_config_file(
"example-security-group.yaml")["provider"]
# a multi-node-type config -- will add head/worker stuff and security group
# info to this.
base_config = helpers.load_aws_example_config_file("example-full.yaml")
config = copy.deepcopy(base_config)
# Add security group data
config["provider"] = provider_data
# Add head and worker fields.
head_node_config = config["available_node_types"]["ray.head.default"][
"node_config"]
worker_node_config = config["available_node_types"]["ray.worker.default"][
"node_config"]
head_node_config["SubnetIds"] = [subnet_id]
worker_node_config["SubnetIds"] = [subnet_id]
# Generate stubs
stubs.configure_iam_role_default(iam_client_stub)
stubs.configure_key_pair_default(ec2_client_stub)
# Only one of these (the one specified in the available_node_types)
# is in the correct vpc.
# This list of subnets is generated by the ec2.subnets.all() call
# and then ignored, since available_node_types already specify
# subnet_ids.
stubs.describe_a_thousand_subnets_in_different_vpcs(ec2_client_stub)
# The rest of the stubbing logic is copied from
# test_create_sg_with_custom_inbound_rules_and_name.
# expect to describe the head subnet ID
stubs.describe_subnets_echo(ec2_client_stub, [DEFAULT_SUBNET])
# given no existing security groups within the VPC...
stubs.describe_no_security_groups(ec2_client_stub)
# expect to create a security group on the head node VPC
stubs.create_sg_echo(ec2_client_stub, DEFAULT_SG_WITH_NAME)
# expect new head security group details to be retrieved after creation
stubs.describe_sgs_on_vpc(
ec2_client_stub,
[DEFAULT_SUBNET["VpcId"]],
[DEFAULT_SG_WITH_NAME],
)
# given custom existing default head security group inbound rules...
# expect to authorize both default and custom inbound rules
stubs.authorize_sg_ingress(
ec2_client_stub,
DEFAULT_SG_WITH_NAME_AND_RULES,
)
# given the prior modification to the head security group...
# expect the next read of a head security group property to reload it
stubs.describe_sg_echo(ec2_client_stub, DEFAULT_SG_WITH_NAME_AND_RULES)
_get_subnets_or_die.cache_clear()
# given our mocks and the config as input...
# expect the config to be validated and bootstrapped successfully
bootstrapped_config = helpers.bootstrap_aws_config(config)
# expect the bootstrapped config to have the custom security group...
# name and in bound rules
assert (bootstrapped_config["provider"]["security_group"]["GroupName"] ==
DEFAULT_SG_WITH_NAME_AND_RULES["GroupName"])
assert (bootstrapped_config["provider"]["security_group"]["IpPermissions"]
== CUSTOM_IN_BOUND_RULES)
# Confirming correct security group got filled for head and workers
sg_id = DEFAULT_SG["GroupId"]
for node_type in bootstrapped_config["available_node_types"].values():
node_config = node_type["node_config"]
assert node_config["SecurityGroupIds"] == [sg_id]
# Confirming boostrap config updates available node types with
# default KeyName
for node_type in bootstrapped_config["available_node_types"].values():
node_config = node_type["node_config"]
assert node_config["KeyName"] == DEFAULT_KEY_PAIR["KeyName"]
# Confirm security group is in the right VPC.
# (Doesn't really confirm anything except for the structure of this test
# data.)
bootstrapped_head_type = bootstrapped_config["head_node_type"]
bootstrapped_types = bootstrapped_config["available_node_types"]
bootstrapped_head_config = bootstrapped_types[bootstrapped_head_type][
"node_config"]
assert DEFAULT_SG["VpcId"] == DEFAULT_SUBNET["VpcId"]
assert DEFAULT_SUBNET["SubnetId"] == bootstrapped_head_config["SubnetIds"][
0]
# ssh private key filled in
assert "ssh_private_key" in bootstrapped_config["auth"]
# expect no pending responses left in IAM or EC2 client stub queues
iam_client_stub.assert_no_pending_responses()
ec2_client_stub.assert_no_pending_responses()
