[Compute] vmss create: add --data-disk-iops and --data-disk-mbps

src/azure-cli/azure/cli/command_modules/vm/_help.py

@@ -2409,7 +2409,10 @@
     text: az vmss update --name MyScaleSet --resource-group MyResourceGroup --instance-id 4 --protect-from-scale-set-actions False --protect-from-scale-in
   - name: Update a VM instance's protection policies.
     text: az vmss update --name MyScaleSet --resource-group MyResourceGroup --instance-id 4 --set protectionPolicy.protectFromScaleIn=True protectionPolicy.protectFromScaleSetActions=False
-
+  - name: Update a VM instance's Read-Write IOPS of the managed disk.
+    text: az vmss update --name MyScaleSet --resource-group MyResourceGroup --set virtualMachineProfile.storageProfile.dataDisks[0].diskIOPSReadWrite=444
+  - name: Update a VM instance's bandwidth in MB per second of the managed disk.
+    text: az vmss update --name MyScaleSet --resource-group MyResourceGroup --set virtualMachineProfile.storageProfile.dataDisks[0].diskMBpsReadWrite=66
 """
 
 helps['vmss update-instances'] = """

src/azure-cli/azure/cli/command_modules/vm/_params.py

@@ -671,6 +671,8 @@ def load_arguments(self, _):
             c.argument('os_disk_encryption_set', min_api='2019-07-01', help='Name or ID of disk encryption set for OS disk.')
             c.argument('data_disk_encryption_sets', nargs='+', min_api='2019-07-01',
                        help='Names or IDs (space delimited) of disk encryption sets for data disks.')
+            c.argument('data_disk_iops', min_api='2019-07-01', nargs='+', type=int, help='Specify the Read-Write IOPS (space delimited) for the managed disk. Should be used only when StorageAccountType is UltraSSD_LRS. If not specified, a default value would be assigned based on diskSizeGB.')
+            c.argument('data_disk_mbps', min_api='2019-07-01', nargs='+', type=int, help='Specify the bandwidth in MB per second (space delimited) for the managed disk. Should be used only when StorageAccountType is UltraSSD_LRS. If not specified, a default value would be assigned based on diskSizeGB.')
 
         with self.argument_context(scope, arg_group='Network') as c:
             c.argument('vnet_name', help='Name of the virtual network when creating a new one or referencing an existing one.')

src/azure-cli/azure/cli/command_modules/vm/_template_builder.py

@@ -659,7 +659,7 @@ def build_vmss_storage_account_pool_resource(_, loop_name, location, tags, stora
     return storage_resource
 
 
-# pylint: disable=too-many-locals, too-many-branches, too-many-statements
+# pylint: disable=too-many-locals, too-many-branches, too-many-statements, too-many-lines
 def build_vmss_resource(cmd, name, naming_prefix, location, tags, overprovision, upgrade_policy_mode,
                         vm_sku, instance_count, ip_config_name, nic_name, subnet_id,
                         public_ip_per_vm, vm_domain_name, dns_servers, nsg, accelerated_networking,
@@ -671,7 +671,8 @@ def build_vmss_resource(cmd, name, naming_prefix, location, tags, overprovision,
                         secrets=None, license_type=None, zones=None, priority=None, eviction_policy=None,
                         application_security_groups=None, ultra_ssd_enabled=None, proximity_placement_group=None,
                         terminate_notification_time=None, max_price=None, scale_in_policy=None,
-                        os_disk_encryption_set=None, data_disk_encryption_sets=None):
+                        os_disk_encryption_set=None, data_disk_encryption_sets=None,
+                        data_disk_iops=None, data_disk_mbps=None):
 
     # Build IP configuration
     ip_configuration = {
@@ -760,6 +761,16 @@ def build_vmss_resource(cmd, name, naming_prefix, location, tags, overprovision,
                 'usage error: Number of --data-disk-encryption-sets mismatches with number of data disks.')
         for i, data_disk in enumerate(data_disks):
             data_disk['managedDisk']['diskEncryptionSet'] = {'id': data_disk_encryption_sets[i]}
+    if data_disk_iops:
+        if len(data_disk_iops) != len(data_disks):
+            raise CLIError('usage error: Number of --data-disk-iops mismatches with number of data disks.')
+        for i, data_disk in enumerate(data_disks):
+            data_disk['diskIOPSReadWrite'] = data_disk_iops[i]
+    if data_disk_mbps:
+        if len(data_disk_mbps) != len(data_disks):
+            raise CLIError('usage error: Number of --data-disk-mbps mismatches with number of data disks.')
+        for i, data_disk in enumerate(data_disks):
+            data_disk['diskMBpsReadWrite'] = data_disk_mbps[i]
     if data_disks:
         storage_properties['dataDisks'] = data_disks
 

src/azure-cli/azure/cli/command_modules/vm/custom.py

@@ -2139,7 +2139,7 @@ def create_vmss(cmd, vmss_name, resource_group_name, image=None,
                 application_security_groups=None, ultra_ssd_enabled=None, ephemeral_os_disk=None,
                 proximity_placement_group=None, aux_subscriptions=None, terminate_notification_time=None,
                 max_price=None, computer_name_prefix=None, orchestration_mode='ScaleSetVM', scale_in_policy=None,
-                os_disk_encryption_set=None, data_disk_encryption_sets=None):
+                os_disk_encryption_set=None, data_disk_encryption_sets=None, data_disk_iops=None, data_disk_mbps=None):
     from azure.cli.core.commands.client_factory import get_subscription_id
     from azure.cli.core.util import random_string, hash_string
     from azure.cli.core.commands.arm import ArmTemplateBuilder
@@ -2367,7 +2367,8 @@ def _get_public_ip_address_allocation(value, sku):
             ultra_ssd_enabled=ultra_ssd_enabled, proximity_placement_group=proximity_placement_group,
             terminate_notification_time=terminate_notification_time, max_price=max_price,
             scale_in_policy=scale_in_policy, os_disk_encryption_set=os_disk_encryption_set,
-            data_disk_encryption_sets=data_disk_encryption_sets)
+            data_disk_encryption_sets=data_disk_encryption_sets, data_disk_iops=data_disk_iops,
+            data_disk_mbps=data_disk_mbps)
 
         vmss_resource['dependsOn'] = vmss_dependencies