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()