[Compute] add VMSS Rolling Upgrade opt

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

@@ -3,7 +3,7 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-# pylint: disable=line-too-long
+# pylint: disable=line-too-long, too-many-lines
 from argcomplete.completers import FilesCompleter
 
 from knack.arguments import CLIArgumentType
@@ -610,6 +610,18 @@ def load_arguments(self, _):
         with self.argument_context(scope) as c:
             c.argument('terminate_notification_time', min_api='2019-03-01',
                        help='Length of time (in minutes, between 5 and 15) a notification to be sent to the VM on the instance metadata server till the VM gets deleted')
+            c.argument('max_batch_instance_percent', type=int, min_api='2020-12-01',
+                       help='The maximum percent of total virtual machine instances that will be upgraded simultaneously by the rolling upgrade in one batch. Default: 20%')
+            c.argument('max_unhealthy_instance_percent', type=int, min_api='2020-12-01',
+                       help='The maximum percentage of the total virtual machine instances in the scale set that can be simultaneously unhealthy. Default: 20%')
+            c.argument('max_unhealthy_upgraded_instance_percent', type=int, min_api='2020-12-01',
+                       help='The maximum percentage of upgraded virtual machine instances that can be found to be in an unhealthy state. Default: 20%')
+            c.argument('pause_time_between_batches', min_api='2020-12-01',
+                       help='The wait time between completing the update for all virtual machines in one batch and starting the next batch. Default: 0 seconds')
+            c.argument('enable_cross_zone_upgrade', arg_type=get_three_state_flag(), min_api='2020-12-01',
+                       help='Set this Boolean property will allow VMSS to ignore AZ boundaries when constructing upgrade batches, and only consider Update Domain and maxBatchInstancePercent to determine the batch size')
+            c.argument('prioritize_unhealthy_instances', arg_type=get_three_state_flag(), min_api='2020-12-01',
+                       help='Set this Boolean property will lead to all unhealthy instances in a scale set getting upgraded before any healthy instances')
 
     for scope, help_prefix in [('vmss update', 'Update the'), ('vmss wait', 'Wait on the')]:
         with self.argument_context(scope) as c:

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

@@ -757,7 +757,10 @@ def build_vmss_resource(cmd, name, naming_prefix, location, tags, overprovision,
                         terminate_notification_time=None, max_price=None, scale_in_policy=None,
                         os_disk_encryption_set=None, data_disk_encryption_sets=None,
                         data_disk_iops=None, data_disk_mbps=None, automatic_repairs_grace_period=None,
-                        specialized=None, os_disk_size_gb=None, encryption_at_host=None, host_group=None):
+                        specialized=None, os_disk_size_gb=None, encryption_at_host=None, host_group=None,
+                        max_batch_instance_percent=None, max_unhealthy_instance_percent=None,
+                        max_unhealthy_upgraded_instance_percent=None, pause_time_between_batches=None,
+                        enable_cross_zone_upgrade=None, prioritize_unhealthy_instances=None):
 
     # Build IP configuration
     ip_configuration = {
@@ -915,7 +918,15 @@ def build_vmss_resource(cmd, name, naming_prefix, location, tags, overprovision,
     vmss_properties = {
         'overprovision': overprovision,
         'upgradePolicy': {
-            'mode': upgrade_policy_mode
+            'mode': upgrade_policy_mode,
+            'rollingUpgradePolicy': {
+                'maxBatchInstancePercent': max_batch_instance_percent,
+                'maxUnhealthyInstancePercent': max_unhealthy_instance_percent,
+                'maxUnhealthyUpgradedInstancePercent': max_unhealthy_upgraded_instance_percent,
+                'pauseTimeBetweenBatches': pause_time_between_batches,
+                'enableCrossZoneUpgrade': enable_cross_zone_upgrade,
+                'prioritizeUnhealthyInstances': prioritize_unhealthy_instances
+            }
         },
         'virtualMachineProfile': {
             'storageProfile': storage_properties,

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

@@ -2419,7 +2419,9 @@ def create_vmss(cmd, vmss_name, resource_group_name, image=None,
                 max_price=None, computer_name_prefix=None, orchestration_mode='Uniform', scale_in_policy=None,
                 os_disk_encryption_set=None, data_disk_encryption_sets=None, data_disk_iops=None, data_disk_mbps=None,
                 automatic_repairs_grace_period=None, specialized=None, os_disk_size_gb=None, encryption_at_host=None,
-                host_group=None):
+                host_group=None, max_batch_instance_percent=None, max_unhealthy_instance_percent=None,
+                max_unhealthy_upgraded_instance_percent=None, pause_time_between_batches=None,
+                enable_cross_zone_upgrade=None, prioritize_unhealthy_instances=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
@@ -2654,7 +2656,11 @@ def _get_public_ip_address_allocation(value, sku):
             data_disk_encryption_sets=data_disk_encryption_sets, data_disk_iops=data_disk_iops,
             data_disk_mbps=data_disk_mbps, automatic_repairs_grace_period=automatic_repairs_grace_period,
             specialized=specialized, os_disk_size_gb=os_disk_size_gb, encryption_at_host=encryption_at_host,
-            host_group=host_group)
+            host_group=host_group, max_batch_instance_percent=max_batch_instance_percent,
+            max_unhealthy_instance_percent=max_unhealthy_instance_percent,
+            max_unhealthy_upgraded_instance_percent=max_unhealthy_upgraded_instance_percent,
+            pause_time_between_batches=pause_time_between_batches, enable_cross_zone_upgrade=enable_cross_zone_upgrade,
+            prioritize_unhealthy_instances=prioritize_unhealthy_instances)
 
         vmss_resource['dependsOn'] = vmss_dependencies
 
@@ -2944,7 +2950,10 @@ def update_vmss(cmd, resource_group_name, name, license_type=None, no_wait=False
                 protect_from_scale_in=None, protect_from_scale_set_actions=None,
                 enable_terminate_notification=None, terminate_notification_time=None, ultra_ssd_enabled=None,
                 scale_in_policy=None, priority=None, max_price=None, proximity_placement_group=None,
-                enable_automatic_repairs=None, automatic_repairs_grace_period=None, **kwargs):
+                enable_automatic_repairs=None, automatic_repairs_grace_period=None, max_batch_instance_percent=None,
+                max_unhealthy_instance_percent=None, max_unhealthy_upgraded_instance_percent=None,
+                pause_time_between_batches=None, enable_cross_zone_upgrade=None, prioritize_unhealthy_instances=None,
+                **kwargs):
     vmss = kwargs['parameters']
     aux_subscriptions = None
     # pylint: disable=too-many-boolean-expressions
@@ -3024,6 +3033,29 @@ def update_vmss(cmd, resource_group_name, name, license_type=None, no_wait=False
     if proximity_placement_group is not None:
         vmss.proximity_placement_group = {'id': proximity_placement_group}
 
+    if max_batch_instance_percent is not None or max_unhealthy_instance_percent is not None \
+            or max_unhealthy_upgraded_instance_percent is not None or pause_time_between_batches is not None \
+            or enable_cross_zone_upgrade is not None or prioritize_unhealthy_instances is not None:
+        if vmss.upgrade_policy is None:
+            vmss.upgrade_policy = {'rolling_upgrade_policy': None}
+        if vmss.upgrade_policy.rolling_upgrade_policy is None:
+            vmss.upgrade_policy.rolling_upgrade_policy = {
+                'maxBatchInstancePercent': max_batch_instance_percent,
+                'maxUnhealthyInstancePercent': max_unhealthy_instance_percent,
+                'maxUnhealthyUpgradedInstancePercent': max_unhealthy_upgraded_instance_percent,
+                'pauseTimeBetweenBatches': pause_time_between_batches,
+                'enableCrossZoneUpgrade': enable_cross_zone_upgrade,
+                'prioritizeUnhealthyInstances': prioritize_unhealthy_instances
+            }
+        else:
+            vmss.upgrade_policy.rolling_upgrade_policy.max_batch_instance_percent = max_batch_instance_percent
+            vmss.upgrade_policy.rolling_upgrade_policy.max_unhealthy_instance_percent = max_unhealthy_instance_percent
+            vmss.upgrade_policy.rolling_upgrade_policy.max_unhealthy_upgraded_instance_percent = \
+                max_unhealthy_upgraded_instance_percent
+            vmss.upgrade_policy.rolling_upgrade_policy.pause_time_between_batches = pause_time_between_batches
+            vmss.upgrade_policy.rolling_upgrade_policy.enable_cross_zone_upgrade = enable_cross_zone_upgrade
+            vmss.upgrade_policy.rolling_upgrade_policy.prioritize_unhealthy_instances = prioritize_unhealthy_instances
+
     return sdk_no_wait(no_wait, client.virtual_machine_scale_sets.begin_create_or_update,
                        resource_group_name, name, **kwargs)
 

src/azure-cli/azure/cli/command_modules/vm/linter_exclusions.yml

@@ -9,6 +9,44 @@ vmss create:
         public_ip_address_allocation:
             rule_exclusions:
             - missing_parameter_help
+        enable_cross_zone_upgrade:
+            rule_exclusions:
+            - option_length_too_long
+        max_batch_instance_percent:
+            rule_exclusions:
+            - option_length_too_long
+        max_unhealthy_instance_percent:
+            rule_exclusions:
+            - option_length_too_long
+        max_unhealthy_upgraded_instance_percent:
+            rule_exclusions:
+            - option_length_too_long
+        pause_time_between_batches:
+            rule_exclusions:
+            - option_length_too_long
+        prioritize_unhealthy_instances:
+            rule_exclusions:
+            - option_length_too_long
+vmss update:
+    parameters:
+        enable_cross_zone_upgrade:
+            rule_exclusions:
+            - option_length_too_long
+        max_batch_instance_percent:
+            rule_exclusions:
+            - option_length_too_long
+        max_unhealthy_instance_percent:
+            rule_exclusions:
+            - option_length_too_long
+        max_unhealthy_upgraded_instance_percent:
+            rule_exclusions:
+            - option_length_too_long
+        pause_time_between_batches:
+            rule_exclusions:
+            - option_length_too_long
+        prioritize_unhealthy_instances:
+            rule_exclusions:
+            - option_length_too_long
 vmss encryption enable:
     parameters:
         key_encryption_algorithm: