Add provisioned-throughput for gce_disk.

PiperOrigin-RevId: 563156317

perfkitbenchmarker/providers/gcp/flags.py

@@ -200,9 +200,13 @@
                     'The GCS bucket to store the preemptible status when '
                     'running on GCP.')
 flags.DEFINE_integer(
-    'gcp_provisioned_iops', 100000,
-    'Iops to provision for pd-extreme. Defaults to the gcloud '
-    'default of 100000.')
+    'gcp_provisioned_iops', None,
+    'Iops to provision for pd-extreme, hyperdisk-extreme, and '
+    'hyperdisk-balanced.')
+flags.DEFINE_integer('gcp_provisioned_throughput', None,
+                     'Provisioned throughput (MB/s) for hyperdisk-throughput '
+                     'and hyperdisk-balanced volumes in GCP. Leave as None to '
+                     'use GCP defaults.')
 flags.DEFINE_boolean('gcp_create_disks_with_vm', True,
                      'Whether to create PD disks at VM creation time. Defaults '
                      'to True.')

perfkitbenchmarker/providers/gcp/gce_disk.py

@@ -47,9 +47,14 @@
     HYPERDISK_EXTREME,
     HYPERDISK_BALANCED,
 ]
-GCE_REMOTE_EXTREME_DISK_TYPES = [
+GCE_DYNAMIC_IOPS_DISK_TYPES = [
     PD_EXTREME,
     HYPERDISK_EXTREME,
+    HYPERDISK_BALANCED,
+]
+GCE_DYNAMIC_THROUGHPUT_DISK_TYPES = [
+    HYPERDISK_BALANCED,
+    HYPERDISK_THROUGHPUT,
 ]
 
 DISK_TYPE = {disk.STANDARD: PD_STANDARD, disk.REMOTE_SSD: PD_SSD}
@@ -200,8 +205,13 @@ def __init__(self,
     self.replica_zones = replica_zones
     self.region = util.GetRegionFromZone(self.zone)
     self.provisioned_iops = None
-    if self.disk_type in GCE_REMOTE_EXTREME_DISK_TYPES:
+    self.provisioned_throughput = None
+    if self.disk_type in GCE_DYNAMIC_IOPS_DISK_TYPES:
       self.provisioned_iops = FLAGS.gcp_provisioned_iops
+      self.metadata['iops'] = self.provisioned_iops
+    if self.disk_type in GCE_DYNAMIC_THROUGHPUT_DISK_TYPES:
+      self.provisioned_throughput = FLAGS.gcp_provisioned_throughput
+      self.metadata['throughput'] = self.provisioned_throughput
 
     disk_metadata = DISK_METADATA[disk_spec.disk_type]
     if self.replica_zones:
@@ -210,22 +220,16 @@ def __init__(self,
     self.metadata.update(DISK_METADATA[disk_spec.disk_type])
     if self.disk_type == disk.LOCAL:
       self.metadata['interface'] = self.interface
-    if (
-        self.provisioned_iops
-        and self.disk_type in GCE_REMOTE_EXTREME_DISK_TYPES
-    ):
-      self.metadata['provisioned_iops'] = self.provisioned_iops
 
   def _Create(self):
     """Creates the disk."""
     cmd = util.GcloudCommand(self, 'compute', 'disks', 'create', self.name)
     cmd.flags['size'] = self.disk_size
     cmd.flags['type'] = self.disk_type
-    if (
-        self.provisioned_iops
-        and self.disk_type in GCE_REMOTE_EXTREME_DISK_TYPES
-    ):
+    if self.provisioned_iops:
       cmd.flags['provisioned-iops'] = self.provisioned_iops
+    if self.provisioned_throughput:
+      cmd.flags['provisioned-throughput'] = self.provisioned_throughput
     cmd.flags['labels'] = util.MakeFormattedDefaultTags()
     if self.image:
       cmd.flags['image'] = self.image

perfkitbenchmarker/providers/gcp/gce_virtual_machine.py

@@ -753,10 +753,18 @@ def _GenerateCreateCommand(self, ssh_keys_path):
             'boot=no',
             'mode=rw',
         ]
-        if disk_spec.disk_type in [
-            gce_disk.PD_EXTREME,
-        ]:
+        if (
+            FLAGS.gcp_provisioned_iops and
+            disk_spec.disk_type in gce_disk.GCE_DYNAMIC_IOPS_DISK_TYPES
+        ):
           pd_args += [f'provisioned-iops={FLAGS.gcp_provisioned_iops}']
+        if (
+            FLAGS.gcp_provisioned_throughput and
+            disk_spec.disk_type in gce_disk.GCE_DYNAMIC_THROUGHPUT_DISK_TYPES
+        ):
+          pd_args += [
+              f'provisioned-throughput={FLAGS.gcp_provisioned_throughput}'
+          ]
         create_disks.append(','.join(pd_args))
     if create_disks:
       cmd.flags['create-disk'] = create_disks