diff --git a/flintrock/ec2.py b/flintrock/ec2.py index 2bbb5f1..e3f18bb 100644 --- a/flintrock/ec2.py +++ b/flintrock/ec2.py @@ -271,6 +271,7 @@ def add_slaves( spot_price: float, spot_request_duration: str, min_root_ebs_size_gb: int, + ebs_volume_type: str, tags: list, assume_yes: bool): security_group_ids = [ @@ -278,6 +279,7 @@ def add_slaves( for group in self.master_instance.security_groups] block_device_mappings = get_ec2_block_device_mappings( min_root_ebs_size_gb=min_root_ebs_size_gb, + ebs_volume_type=ebs_volume_type, ami=self.master_instance.image_id, region=self.region) availability_zone = self.master_instance.placement['AvailabilityZone'] @@ -653,6 +655,7 @@ def get_or_create_flintrock_security_groups( def get_ec2_block_device_mappings( *, min_root_ebs_size_gb: int, + ebs_volume_type: str, ami: str, region: str) -> 'List[dict]': """ @@ -688,7 +691,7 @@ def get_ec2_block_device_mappings( # of a root instance store volume. 'VolumeSize': min_root_ebs_size_gb, # gp2 is general-purpose SSD - 'VolumeType': 'gp2'}) + 'VolumeType': ebs_volume_type}) del root_device['Ebs']['Encrypted'] block_device_mappings.append(root_device) @@ -853,6 +856,7 @@ def launch( spot_price=None, spot_request_duration=None, min_root_ebs_size_gb, + ebs_volume_type, vpc_id, subnet_id, instance_profile_name, @@ -903,6 +907,7 @@ def launch( security_group_ids = [sg.id for sg in user_security_groups + flintrock_security_groups] block_device_mappings = get_ec2_block_device_mappings( min_root_ebs_size_gb=min_root_ebs_size_gb, + ebs_volume_type=ebs_volume_type, ami=ami, region=region) diff --git a/flintrock/flintrock.py b/flintrock/flintrock.py index 9c0fe67..8d1e770 100644 --- a/flintrock/flintrock.py +++ b/flintrock/flintrock.py @@ -348,6 +348,7 @@ def cli(cli_context, config, provider, debug): @click.option('--ec2-spot-request-duration', default='7d', help="Duration a spot request is valid (e.g. 3d 2h 1m).") @click.option('--ec2-min-root-ebs-size-gb', type=int, default=30) +@click.option('--ec2-ebs-volume-type', default='gp2') @click.option('--ec2-vpc-id', default='', help="Leave empty for default VPC.") @click.option('--ec2-subnet-id', default='') @click.option('--ec2-instance-profile-name', default='') @@ -401,6 +402,7 @@ def launch( ec2_spot_price, ec2_spot_request_duration, ec2_min_root_ebs_size_gb, + ec2_ebs_volume_type, ec2_vpc_id, ec2_subnet_id, ec2_instance_profile_name, @@ -508,6 +510,7 @@ def launch( spot_price=ec2_spot_price, spot_request_duration=ec2_spot_request_duration, min_root_ebs_size_gb=ec2_min_root_ebs_size_gb, + ebs_volume_type=ec2_ebs_volume_type, vpc_id=ec2_vpc_id, subnet_id=ec2_subnet_id, instance_profile_name=ec2_instance_profile_name, @@ -785,6 +788,7 @@ def stop(cli_context, cluster_name, ec2_region, ec2_vpc_id, assume_yes): @click.option('--ec2-spot-request-duration', default='7d', help="Duration a spot request is valid (e.g. 3d 2h 1m).") @click.option('--ec2-min-root-ebs-size-gb', type=int, default=30) +@click.option('--ec2-ebs-volume-type', default='gp2') @click.option('--assume-yes/--no-assume-yes', default=False) @click.option('--ec2-tag', 'ec2_tags', callback=ec2.cli_validate_tags, @@ -803,6 +807,7 @@ def add_slaves( ec2_spot_price, ec2_spot_request_duration, ec2_min_root_ebs_size_gb, + ec2_ebs_volume_type, ec2_tags, assume_yes): """ @@ -831,6 +836,7 @@ def add_slaves( identity_file = ec2_identity_file provider_options = { 'min_root_ebs_size_gb': ec2_min_root_ebs_size_gb, + 'ebs_volume_type': ec2_ebs_volume_type, 'spot_price': ec2_spot_price, 'spot_request_duration': ec2_spot_request_duration, 'tags': ec2_tags