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_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_subnet_given_head_and_worker_sg(iam_client_stub, ec2_client_stub):
stubs.configure_iam_role_default(iam_client_stub)
stubs.configure_key_pair_default(ec2_client_stub)
# list a security group and a thousand subnets in different vpcs
stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG)
stubs.describe_a_thousand_subnets_in_different_vpcs(ec2_client_stub)
config = helpers.bootstrap_aws_example_config_file(
"example-head-and-worker-security-group.yaml")
# check that just the single subnet in the right vpc is filled
assert config["head_node"]["SubnetIds"] == [DEFAULT_SUBNET["SubnetId"]]
assert config["worker_nodes"]["SubnetIds"] == [DEFAULT_SUBNET["SubnetId"]]
# 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()
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()
