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_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_configure_subnet_with_network_interfaces_missing_subnet_errors(
ec2_client_stub, node_type):
empty = {}
default = {"SubnetId": DEFAULT_SUBNET["SubnetId"]}
config = {
"provider": {
"type": "aws",
"region": "us-west-2"
},
"head_node": {
"NetworkInterfaces": [
empty if node_type == "head_node" else default
]
},
"worker_nodes": {
"NetworkInterfaces": [
empty if node_type == "worker_nodes" else default
]
}
}
stubs.configure_subnet_default(ec2_client_stub)
with pytest.raises(Exception) as x:
_ = _configure_subnet(config)
assert "missing a subnet" in str(x.value)
def test_launch_templates(ec2_client_stub, ec2_client_stub_fail_fast,
ec2_client_stub_max_retries):
# given the launch template associated with our default head node type...
# expect to first describe the default launch template by ID
stubs.describe_launch_template_versions_by_id_default(
ec2_client_stub, ["$Latest"])
# given the launch template associated with our default worker node type...
# expect to next describe the same default launch template by name
stubs.describe_launch_template_versions_by_name_default(
ec2_client_stub, ["2"])
# use default stubs to skip ahead to subnet configuration
stubs.configure_key_pair_default(ec2_client_stub)
# given the security groups associated with our launch template...
sgids = [DEFAULT_SG["GroupId"]]
security_groups = [DEFAULT_SG]
# expect to describe all security groups to ensure they share the same VPC
stubs.describe_sgs_by_id(ec2_client_stub, sgids, security_groups)
# use a default stub to skip subnet configuration
stubs.configure_subnet_default(ec2_client_stub)
# given our mocks and an example config file as input...
# expect the config to be loaded, validated, and bootstrapped successfully
config = helpers.bootstrap_aws_example_config_file(
"example-launch-templates.yaml")
# instantiate a new node provider
new_provider = _get_node_provider(
config["provider"],
DEFAULT_CLUSTER_NAME,
False,
)
max_count = 1
for name, node_type in config["available_node_types"].items():
# given our bootstrapped node config as input to create a new node...
# expect to first describe all stopped instances that could be reused
stubs.describe_instances_with_any_filter_consumer(
ec2_client_stub_max_retries)
# given no stopped EC2 instances to reuse...
# expect to create new nodes with the given launch template config
node_cfg = node_type["node_config"]
stubs.run_instances_with_launch_template_consumer(
ec2_client_stub_fail_fast,
config,
node_cfg,
name,
DEFAULT_LT["LaunchTemplateData"],
max_count,
)
tags = helpers.node_provider_tags(config, name)
new_provider.create_node(node_cfg, tags, max_count)
ec2_client_stub.assert_no_pending_responses()
ec2_client_stub_fail_fast.assert_no_pending_responses()
ec2_client_stub_max_retries.assert_no_pending_responses()
def test_network_interfaces(
ec2_client_stub,
iam_client_stub,
ec2_client_stub_fail_fast,
ec2_client_stub_max_retries,
):
# use default stubs to skip ahead to subnet configuration
stubs.configure_iam_role_default(iam_client_stub)
stubs.configure_key_pair_default(ec2_client_stub)
# given the security groups associated with our network interfaces...
sgids = ["sg-00000000", "sg-11111111", "sg-22222222", "sg-33333333"]
security_groups = []
suffix = 0
for sgid in sgids:
sg = copy.deepcopy(DEFAULT_SG)
sg["GroupName"] += f"-{suffix}"
sg["GroupId"] = sgid
security_groups.append(sg)
suffix += 1
# expect to describe all security groups to ensure they share the same VPC
stubs.describe_sgs_by_id(ec2_client_stub, sgids, security_groups)
# use a default stub to skip subnet configuration
stubs.configure_subnet_default(ec2_client_stub)
# given our mocks and an example config file as input...
# expect the config to be loaded, validated, and bootstrapped successfully
config = helpers.bootstrap_aws_example_config_file(
"example-network-interfaces.yaml")
# instantiate a new node provider
new_provider = _get_node_provider(
config["provider"],
DEFAULT_CLUSTER_NAME,
False,
)
for name, node_type in config["available_node_types"].items():
node_cfg = node_type["node_config"]
tags = helpers.node_provider_tags(config, name)
# given our bootstrapped node config as input to create a new node...
# expect to first describe all stopped instances that could be reused
stubs.describe_instances_with_any_filter_consumer(
ec2_client_stub_max_retries)
# given no stopped EC2 instances to reuse...
# expect to create new nodes with the given network interface config
stubs.run_instances_with_network_interfaces_consumer(
ec2_client_stub_fail_fast,
node_cfg["NetworkInterfaces"],
)
new_provider.create_node(node_cfg, tags, 1)
iam_client_stub.assert_no_pending_responses()
ec2_client_stub.assert_no_pending_responses()
ec2_client_stub_fail_fast.assert_no_pending_responses()
ec2_client_stub_max_retries.assert_no_pending_responses()
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_configure_subnet_with_network_interfaces_with_subnet(ec2_client_stub):
config = {
"provider": {
"type": "aws",
"region": "us-west-2"
},
"head_node": {
"NetworkInterfaces": [{
"SubnetId": DEFAULT_SUBNET["SubnetId"]
}]
},
"worker_nodes": {
"NetworkInterfaces": [{
"SubnetId": DEFAULT_SUBNET["SubnetId"]
}]
}
}
stubs.configure_subnet_default(ec2_client_stub)
expected = copy.deepcopy(config)
config = _configure_subnet(config)
assert expected == config
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()