allow ebs volume type change using configuration, hardcoded gp2 break…

… gp3 image launch
nchammas · Sep 18, 2023 · b96d623 · b96d623
1 parent 01f8ef9
commit b96d623
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/flintrock/ec2.py b/flintrock/ec2.py
@@ -271,13 +271,15 @@ 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 = [
             group['GroupId']
             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
@@ -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