Aws: add support for volume iops and throughput

Users can request specific IOPS and throughput allocations from EC2.
The availability and defaults vary for volume type, but IOPS are
available for all volumes, and throughput is available on gp3 volumes.

Change-Id: Icc7432d8ce1c3514bfe9d8fda20bd399b67ede7a

doc/source/aws.rst

@@ -411,6 +411,20 @@ Selecting the ``aws`` driver adds the following options to the
          omitted, the volume size reported for the imported snapshot
          will be used.
 
+      .. attr:: iops
+         :type: int
+
+         The number of I/O operations per second to be provisioned for
+         the volume.  The default varies based on the volume type; see
+         the documentation under `EBS volume type`_ for the specific
+         volume type for details.
+
+      .. attr:: throughput
+         :type: int
+
+         The throughput of the volume in MiB/s.  This is only valid for
+         ``gp3`` volumes.
+
       .. attr:: tags
          :type: dict
          :default: None
@@ -626,6 +640,20 @@ Selecting the ``aws`` driver adds the following options to the
 
               If given, the size of the root EBS volume, in GiB.
 
+           .. attr:: iops
+              :type: int
+
+              The number of I/O operations per second to be
+              provisioned for the volume.  The default varies based on
+              the volume type; see the documentation under `EBS volume
+              type`_ for the specific volume type for details.
+
+           .. attr:: throughput
+              :type: int
+
+              The throughput of the volume in MiB/s.  This is only
+              valid for ``gp3`` volumes.
+
            .. attr:: userdata
               :type: str
               :default: None

nodepool/driver/aws/adapter.py

@@ -456,20 +456,24 @@ class AwsAdapter(statemachine.Adapter):
         volume_size = provider_image.volume_size or snap.volume_size
         # Register the snapshot as an AMI
         with self.rate_limiter:
+            bdm = {
+                'DeviceName': '/dev/sda1',
+                'Ebs': {
+                    'DeleteOnTermination': True,
+                    'SnapshotId': task[
+                        'SnapshotTaskDetail']['SnapshotId'],
+                    'VolumeSize': volume_size,
+                    'VolumeType': provider_image.volume_type,
+                },
+            }
+            if provider_image.iops:
+                bdm['Ebs']['Iops'] = provider_image.iops
+            if provider_image.throughput:
+                bdm['Ebs']['Throughput'] = provider_image.throughput
+
             register_response = self.ec2_client.register_image(
                 Architecture=provider_image.architecture,
-                BlockDeviceMappings=[
-                    {
-                        'DeviceName': '/dev/sda1',
-                        'Ebs': {
-                            'DeleteOnTermination': True,
-                            'SnapshotId': task[
-                                'SnapshotTaskDetail']['SnapshotId'],
-                            'VolumeSize': volume_size,
-                            'VolumeType': provider_image.volume_type,
-                        },
-                    },
-                ],
+                BlockDeviceMappings=[bdm],
                 RootDeviceName='/dev/sda1',
                 VirtualizationType='hvm',
                 EnaSupport=provider_image.ena_support,
@@ -817,6 +821,10 @@ class AwsAdapter(statemachine.Adapter):
                     mapping['Ebs']['VolumeSize'] = label.volume_size
                 if label.volume_type:
                     mapping['Ebs']['VolumeType'] = label.volume_type
+                if label.iops:
+                    mapping['Ebs']['Iops'] = label.iops
+                if label.throughput:
+                    mapping['Ebs']['Throughput'] = label.throughput
                 # If the AMI is a snapshot, we cannot supply an "encrypted"
                 # parameter
                 if 'Encrypted' in mapping['Ebs']:

nodepool/driver/aws/config.py

@@ -104,6 +104,8 @@ class AwsProviderDiskImage(ConfigValue):
         self.ena_support = image.get('ena-support', True)
         self.volume_size = image.get('volume-size', None)
         self.volume_type = image.get('volume-type', 'gp2')
+        self.iops = image.get('iops', None)
+        self.throughput = image.get('throughput', None)
 
     @property
     def external_name(self):
@@ -124,6 +126,8 @@ class AwsProviderDiskImage(ConfigValue):
             'ena-support': bool,
             'volume-size': int,
             'volume-type': str,
+            'iops': int,
+            'throughput': int,
             'tags': dict,
         }
 
@@ -166,6 +170,8 @@ class AwsLabel(ConfigValue):
         self.key_name = label.get('key-name')
         self.volume_type = label.get('volume-type')
         self.volume_size = label.get('volume-size')
+        self.iops = label.get('iops', None)
+        self.throughput = label.get('throughput', None)
         self.userdata = label.get('userdata', None)
         self.iam_instance_profile = label.get('iam-instance-profile', None)
         self.tags = label.get('tags', {})
@@ -182,6 +188,8 @@ class AwsLabel(ConfigValue):
             'ebs-optimized': bool,
             'volume-type': str,
             'volume-size': int,
+            'iops': int,
+            'throughput': int,
             'userdata': str,
             'iam-instance-profile': {
                 v.Exclusive('name', 'iam_instance_profile_id'): str,

nodepool/tests/fixtures/aws/diskimage.yaml

@@ -31,6 +31,9 @@ providers:
       - name: fake-image
         tags:
           provider_metadata: provider
+        volume-type: gp3
+        iops: 1000
+        throughput: 100
     pools:
       - name: main
         max-servers: 1
@@ -44,6 +47,8 @@ providers:
             diskimage: fake-image
             instance-type: t3.medium
             key-name: zuul
+            iops: 2000
+            throughput: 200
 
 diskimages:
   - name: fake-image

nodepool/tests/unit/test_driver_aws.py

@@ -609,6 +609,12 @@ class TestDriverAws(tests.DBTestCase):
         self.assertEqual(node.shell_type, None)
         self.assertEqual(node.attributes,
                          {'key1': 'value1', 'key2': 'value2'})
+        self.assertEqual(
+            self.create_instance_calls[0]['BlockDeviceMappings'][0]['Ebs']
+            ['Iops'], 2000)
+        self.assertEqual(
+            self.create_instance_calls[0]['BlockDeviceMappings'][0]['Ebs']
+            ['Throughput'], 200)
 
     def test_aws_diskimage_removal(self):
         configfile = self.setup_config('aws/diskimage.yaml')

releasenotes/notes/aws-iops-6f6f54f0b111c13b.yaml

@@ -0,0 +1,8 @@
+---
+features:
+  - |
+    The AWS driver now support specifying volume IOPS and throughput; see:
+    :attr:`providers.[aws].pools.labels.iops`,
+    :attr:`providers.[aws].pools.labels.throughput`,
+    :attr:`providers.[aws].diskimages.iops`, and
+    :attr:`providers.[aws].diskimages.throughput`.