[Compute] az vmss create/update: Add new parameters --enable-automatic-zone-balancing to support setting automatic zone rebalancing policy

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

@@ -863,6 +863,9 @@ def load_arguments(self, _):
             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.')
             c.argument('security_type', arg_type=get_enum_type(["TrustedLaunch", "Standard", "ConfidentialVM"], default=None),
                        help='Specify the security type of the virtual machine scale set. The value Standard can be used if subscription has feature flag UseStandardSecurityType registered under Microsoft.Compute namespace. Refer to https://learn.microsoft.com/en-us/azure/azure-resource-manager/management/preview-features for steps to enable required feature.')
+            c.argument('enable_automatic_zone_balancing', arg_type=get_three_state_flag(), options_list=['--enable-automatic-zone-balancing', '--enable-zone-balancing'], min_api='2024-11-01', help='Specify whether automatic AZ balancing should be enabled on the virtualmachine scale set.')
+            c.argument('automatic_zone_balancing_strategy', arg_type=get_enum_type(self.get_models('RebalanceStrategy')), options_list=['--automatic-zone-balancing-strategy', '--balancing-strategy'], min_api='2024-11-01', help='Type of rebalance strategy that will be used for rebalancing virtualmachines in the scale set across availability zones.')
+            c.argument('automatic_zone_balancing_behavior', arg_type=get_enum_type(self.get_models('RebalanceBehavior')), options_list=['--automatic-zone-balancing-behavior', '--balancing-behavior'], min_api='2024-11-01', help='Type of rebalance behavior that will be used for recreating virtualmachines in the scale set across availability zones.')
 
     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

@@ -1042,7 +1042,8 @@ def build_vmss_resource(cmd, name, computer_name_prefix, location, tags, overpro
                         enable_user_redeploy_scheduled_events=None, skuprofile_vmsizes=None, skuprofile_allostrat=None,
                         security_posture_reference_is_overridable=None, zone_balance=None, wire_server_mode=None,
                         imds_mode=None, wire_server_access_control_profile_reference_id=None,
-                        imds_access_control_profile_reference_id=None):
+                        imds_access_control_profile_reference_id=None, enable_automatic_zone_balancing=None,
+                        automatic_zone_balancing_strategy=None, automatic_zone_balancing_behavior=None):
 
     # Build IP configuration
     ip_configuration = {}
@@ -1530,12 +1531,26 @@ def build_vmss_resource(cmd, name, computer_name_prefix, location, tags, overpro
     if scale_in_policy:
         vmss_properties['scaleInPolicy'] = {'rules': scale_in_policy}
 
-    if enable_resilient_vm_creation is not None or enable_resilient_vm_deletion is not None:
-        resiliency_policy = {}
-        if enable_resilient_vm_creation is not None:
-            resiliency_policy['resilientVMCreationPolicy'] = {'enabled': enable_resilient_vm_creation}
-        if enable_resilient_vm_deletion is not None:
-            resiliency_policy['resilientVMDeletionPolicy'] = {'enabled': enable_resilient_vm_deletion}
+    resiliency_policy = {}
+    if enable_resilient_vm_creation is not None:
+        resiliency_policy['resilientVMCreationPolicy'] = {'enabled': enable_resilient_vm_creation}
+    if enable_resilient_vm_deletion is not None:
+        resiliency_policy['resilientVMDeletionPolicy'] = {'enabled': enable_resilient_vm_deletion}
+
+    automatic_zone_rebalancing_policy = {}
+    if enable_automatic_zone_balancing is not None:
+        automatic_zone_rebalancing_policy['enabled'] = enable_automatic_zone_balancing
+
+    if automatic_zone_balancing_strategy is not None:
+        automatic_zone_rebalancing_policy['rebalanceStrategy'] = automatic_zone_balancing_strategy
+
+    if automatic_zone_balancing_behavior is not None:
+        automatic_zone_rebalancing_policy['rebalanceBehavior'] = automatic_zone_balancing_behavior
+
+    if automatic_zone_rebalancing_policy:
+        resiliency_policy['automaticZoneRebalancingPolicy'] = automatic_zone_rebalancing_policy
+
+    if resiliency_policy:
         vmss_properties['resiliencyPolicy'] = resiliency_policy
 
     security_profile = {}

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

@@ -3246,7 +3246,8 @@ def create_vmss(cmd, vmss_name, resource_group_name, image=None,
                 enable_user_redeploy_scheduled_events=None, skuprofile_vmsizes=None, skuprofile_allostrat=None,
                 security_posture_reference_is_overridable=None, zone_balance=None, wire_server_mode=None,
                 imds_mode=None, wire_server_access_control_profile_reference_id=None,
-                imds_access_control_profile_reference_id=None):
+                imds_access_control_profile_reference_id=None, enable_automatic_zone_balancing=None,
+                automatic_zone_balancing_strategy=None, automatic_zone_balancing_behavior=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
@@ -3564,7 +3565,10 @@ def _get_public_ip_address_allocation(value, sku):
             security_posture_reference_is_overridable=security_posture_reference_is_overridable,
             zone_balance=zone_balance, wire_server_mode=wire_server_mode, imds_mode=imds_mode,
             wire_server_access_control_profile_reference_id=wire_server_access_control_profile_reference_id,
-            imds_access_control_profile_reference_id=imds_access_control_profile_reference_id)
+            imds_access_control_profile_reference_id=imds_access_control_profile_reference_id,
+            enable_automatic_zone_balancing=enable_automatic_zone_balancing,
+            automatic_zone_balancing_strategy=automatic_zone_balancing_strategy,
+            automatic_zone_balancing_behavior=automatic_zone_balancing_behavior)
 
         vmss_resource['dependsOn'] = vmss_dependencies
 
@@ -4015,7 +4019,8 @@ def update_vmss(cmd, resource_group_name, name, license_type=None, no_wait=False
                 upgrade_policy_mode=None, enable_auto_os_upgrade=None, skuprofile_vmsizes=None,
                 skuprofile_allostrat=None, security_posture_reference_is_overridable=None, zone_balance=None,
                 wire_server_mode=None, imds_mode=None, wire_server_access_control_profile_reference_id=None,
-                imds_access_control_profile_reference_id=None, **kwargs):
+                imds_access_control_profile_reference_id=None, enable_automatic_zone_balancing=None,
+                automatic_zone_balancing_strategy=None, automatic_zone_balancing_behavior=None, **kwargs):
     vmss = kwargs['parameters']
     aux_subscriptions = None
     # pylint: disable=too-many-boolean-expressions
@@ -4332,6 +4337,28 @@ def update_vmss(cmd, resource_group_name, name, license_type=None, no_wait=False
         if enable_resilient_deletion is not None:
             resiliency_policy.resilient_vm_deletion_policy = {'enabled': enable_resilient_deletion}
 
+    if enable_automatic_zone_balancing is not None or automatic_zone_balancing_strategy is not None or \
+            automatic_zone_balancing_behavior is not None:
+        resiliency_policy = vmss.resiliency_policy
+        if resiliency_policy is None:
+            ResiliencyPolicy = cmd.get_models('ResiliencyPolicy')
+            AutomaticZoneRebalancingPolicy = cmd.get_models('AutomaticZoneRebalancingPolicy')
+            resiliency_policy = ResiliencyPolicy()
+            resiliency_policy.automatic_zone_rebalancing_policy = AutomaticZoneRebalancingPolicy()
+        elif resiliency_policy.automatic_zone_rebalancing_policy is None:
+            AutomaticZoneRebalancingPolicy = cmd.get_models('AutomaticZoneRebalancingPolicy')
+            resiliency_policy.automatic_zone_rebalancing_policy = AutomaticZoneRebalancingPolicy()
+
+        if enable_automatic_zone_balancing is not None:
+            resiliency_policy.automatic_zone_rebalancing_policy.enabled = enable_automatic_zone_balancing
+
+        if automatic_zone_balancing_strategy is not None:
+            resiliency_policy.automatic_zone_rebalancing_policy.rebalance_strategy = automatic_zone_balancing_strategy
+
+        if automatic_zone_balancing_behavior is not None:
+            resiliency_policy.automatic_zone_rebalancing_policy.rebalance_behavior = automatic_zone_balancing_behavior
+        vmss.resiliency_policy = resiliency_policy
+
     if zones is not None:
         vmss.zones = zones