Merge "Aws: add support for volume iops and throughput"

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,10 +456,7 @@ class AwsAdapter(statemachine.Adapter):
         volume_size = provider_image.volume_size or snap.volume_size
         # Register the snapshot as an AMI
         with self.rate_limiter:
-            register_response = self.ec2_client.register_image(
+            bdm = {
-                Architecture=provider_image.architecture,
-                BlockDeviceMappings=[
-                    {
                 'DeviceName': '/dev/sda1',
                 'Ebs': {
                     'DeleteOnTermination': True,
@@ -468,8 +465,15 @@ class AwsAdapter(statemachine.Adapter):
                     '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=[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`.