[Compute] az vmss create/update: Add new paramter --zone-balance to support setting zone balance (#30692)

Author: Jing-song

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

@@ -871,6 +871,7 @@ def load_arguments(self, _):
                        help='List of virtual machine extensions to exclude when applying the Security Posture. Either a Json string or a file path is acceptable. '
                             'Please refer to https://docs.microsoft.com/rest/api/compute/virtualmachinescalesets/get#virtualmachineextension for the data format.')
             c.argument('security_posture_reference_is_overridable', arg_type=get_three_state_flag(), min_api='2024-03-01', options_list=['--security-posture-reference-is-overridable', '--is-overridable'], help='Whether the security posture can be overridden by the user.')
+            c.argument('zone_balance', arg_type=get_three_state_flag(), min_api='2017-12-01', help='Whether to force strictly even Virtual Machine distribution cross x-zones in case there is zone outage.')
 
     with self.argument_context('vmss update') as c:
         c.argument('instance_id', id_part='child_name_1', help="Update the VM instance with this ID. If missing, update the VMSS.")

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

@@ -1004,7 +1004,7 @@ def build_vmss_resource(cmd, name, computer_name_prefix, location, tags, overpro
                         enable_resilient_vm_creation=None, enable_resilient_vm_deletion=None,
                         additional_scheduled_events=None, enable_user_reboot_scheduled_events=None,
                         enable_user_redeploy_scheduled_events=None, skuprofile_vmsizes=None, skuprofile_allostrat=None,
-                        security_posture_reference_is_overridable=None):
+                        security_posture_reference_is_overridable=None, zone_balance=None):
 
     # Build IP configuration
     ip_configuration = {}
@@ -1436,6 +1436,9 @@ def build_vmss_resource(cmd, name, computer_name_prefix, location, tags, overpro
     if proximity_placement_group:
         vmss_properties['proximityPlacementGroup'] = {'id': proximity_placement_group}
 
+    if zone_balance is not None:
+        vmss_properties['zoneBalance'] = zone_balance
+
     scheduled_events_profile = {}
     if terminate_notification_time is not None:
         scheduled_events_profile.update({

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

@@ -3234,7 +3234,7 @@ def create_vmss(cmd, vmss_name, resource_group_name, image=None,
                 enable_resilient_creation=None, enable_resilient_deletion=None,
                 additional_scheduled_events=None, enable_user_reboot_scheduled_events=None,
                 enable_user_redeploy_scheduled_events=None, skuprofile_vmsizes=None, skuprofile_allostrat=None,
-                security_posture_reference_is_overridable=None):
+                security_posture_reference_is_overridable=None, zone_balance=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
@@ -3549,7 +3549,8 @@ def _get_public_ip_address_allocation(value, sku):
             enable_user_reboot_scheduled_events=enable_user_reboot_scheduled_events,
             enable_user_redeploy_scheduled_events=enable_user_redeploy_scheduled_events,
             skuprofile_vmsizes=skuprofile_vmsizes, skuprofile_allostrat=skuprofile_allostrat,
-            security_posture_reference_is_overridable=security_posture_reference_is_overridable)
+            security_posture_reference_is_overridable=security_posture_reference_is_overridable,
+            zone_balance=zone_balance)
 
         vmss_resource['dependsOn'] = vmss_dependencies
 
@@ -3989,7 +3990,7 @@ def update_vmss(cmd, resource_group_name, name, license_type=None, no_wait=False
                 ephemeral_os_disk=None, ephemeral_os_disk_option=None, zones=None, additional_scheduled_events=None,
                 enable_user_reboot_scheduled_events=None, enable_user_redeploy_scheduled_events=None,
                 upgrade_policy_mode=None, enable_auto_os_upgrade=None, skuprofile_vmsizes=None,
-                skuprofile_allostrat=None, security_posture_reference_is_overridable=None, **kwargs):
+                skuprofile_allostrat=None, security_posture_reference_is_overridable=None, zone_balance=None, **kwargs):
     vmss = kwargs['parameters']
     aux_subscriptions = None
     # pylint: disable=too-many-boolean-expressions
@@ -4293,6 +4294,9 @@ def update_vmss(cmd, resource_group_name, name, license_type=None, no_wait=False
     if zones is not None:
         vmss.zones = zones
 
+    if zone_balance is not None:
+        vmss.zone_balance = zone_balance
+
     return sdk_no_wait(no_wait, client.virtual_machine_scale_sets.begin_create_or_update,
                        resource_group_name, name, **kwargs)