[AKS] az aks nodepool update: Add --zones for automatic zone placement

src/azure-cli/azure/cli/command_modules/acs/_consts.py

@@ -6,6 +6,8 @@
 from enum import Enum
 
 # consts for AgentPool
+CONST_AGENTPOOL_UPDATE_ZONES_API_VERSION = "2026-01-02-preview"
+
 # priority
 CONST_SCALE_SET_PRIORITY_REGULAR = "Regular"
 CONST_SCALE_SET_PRIORITY_SPOT = "Spot"

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

@@ -2210,6 +2210,9 @@
   - name: --mode
     type: string
     short-summary: The mode for a node pool which defines a node pool's primary function. If set as "System", AKS prefers system pods scheduling to node pools with mode `System`. Learn more at https://aka.ms/aks/nodepool/mode.
+  - name: --zones -z
+    type: string array
+    short-summary: Use `auto` to migrate a regional node pool to automatic zone placement. Other availability zone changes are subject to service restrictions.
   - name: --max-surge
     type: string
     short-summary: Extra nodes used to speed upgrade. When specified, it represents the number or percent used, eg. 5 or 33%
@@ -2288,6 +2291,8 @@
     text: az aks nodepool update --disable-cluster-autoscaler -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster
   - name: Update min-count or max-count for cluster autoscaler.
     text: az aks nodepool update --update-cluster-autoscaler --min-count 1 --max-count 10 -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster
+  - name: Migrate a regional node pool to automatic zone placement.
+    text: az aks nodepool update -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster --zones auto
 """
 
 helps["aks nodepool upgrade"] = """

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

@@ -1192,6 +1192,14 @@ def load_arguments(self, _):
         c.argument("undrainable_node_behavior")
         c.argument('mode', get_enum_type(node_mode_types))
         c.argument('scale_down_mode', arg_type=get_enum_type(scale_down_modes))
+        c.argument(
+            'zones',
+            zones_type,
+            options_list=['--zones', '-z'],
+            is_preview=True,
+            help='Use "auto" to migrate a regional node pool to automatic zone placement. '
+                 'Other availability zone changes are subject to service restrictions.',
+        )
         c.argument('allowed_host_ports', nargs='+', validator=validate_allowed_host_ports)
         c.argument('asg_ids', nargs='+', validator=validate_application_security_groups)
         c.argument('os_sku', arg_type=get_enum_type(node_os_skus_update), validator=validate_os_sku)

src/azure-cli/azure/cli/command_modules/acs/agentpool_decorator.py

@@ -11,6 +11,7 @@
 
 from azure.cli.command_modules.acs._client_factory import cf_agent_pools
 from azure.cli.command_modules.acs._consts import (
+    CONST_AGENTPOOL_UPDATE_ZONES_API_VERSION,
     CONST_AVAILABILITY_SET,
     CONST_DEFAULT_NODE_OS_TYPE,
     CONST_DEFAULT_NODE_VM_SIZE,
@@ -48,6 +49,7 @@
 )
 from azure.cli.core.cloud import get_active_cloud
 from azure.cli.core.commands import AzCliCommand
+from azure.cli.core.commands.client_factory import get_mgmt_service_client
 from azure.cli.core.profiles import ResourceType
 from azure.cli.core.util import get_file_json, sdk_no_wait, read_file_content
 from knack.log import get_logger
@@ -2861,6 +2863,18 @@ def update_vm_properties(self, agentpool: AgentPool) -> AgentPool:
             agentpool.mode = mode
         return agentpool
 
+    def update_zones(self, agentpool: AgentPool) -> AgentPool:
+        """Update availability zones for the AgentPool object.
+
+        :return: the AgentPool object
+        """
+        self._ensure_agentpool(agentpool)
+
+        zones = self.context.raw_param.get("zones")
+        if zones is not None:
+            agentpool.availability_zones = zones
+        return agentpool
+
     def update_network_profile(self, agentpool: AgentPool) -> AgentPool:
         self._ensure_agentpool(agentpool)
 
@@ -2982,6 +2996,8 @@ def update_agentpool_profile_default(self, agentpools: List[AgentPool] = None) -
         agentpool = self.update_upgrade_settings(agentpool)
         # update misc vm properties
         agentpool = self.update_vm_properties(agentpool)
+        # update availability zones
+        agentpool = self.update_zones(agentpool)
         # update network profile
         agentpool = self.update_network_profile(agentpool)
         # update os sku
@@ -3010,11 +3026,24 @@ def update_agentpool(self, agentpool: AgentPool) -> AgentPool:
         """
         self._ensure_agentpool(agentpool)
 
+        update_client = self.client
+        # Zone migration is not accepted by the stable API yet, so scope the preview client to this PUT.
+        if (
+            self.resource_type == ResourceType.MGMT_CONTAINERSERVICE and
+            self.agentpool_decorator_mode == AgentPoolDecoratorMode.STANDALONE and
+            self.context.raw_param.get("zones") is not None
+        ):
+            update_client = get_mgmt_service_client(
+                self.cmd.cli_ctx,
+                ResourceType.MGMT_CONTAINERSERVICE,
+                api_version=CONST_AGENTPOOL_UPDATE_ZONES_API_VERSION,
+            ).agent_pools
+
         active_cloud = get_active_cloud(self.cmd.cli_ctx)
         if active_cloud.profile != "latest":
             return sdk_no_wait(
                 self.context.get_no_wait(),
-                self.client.begin_create_or_update,
+                update_client.begin_create_or_update,
                 self.context.get_resource_group_name(),
                 self.context.get_cluster_name(),
                 self.context.get_nodepool_name(),
@@ -3024,7 +3053,7 @@ def update_agentpool(self, agentpool: AgentPool) -> AgentPool:
 
         return sdk_no_wait(
             self.context.get_no_wait(),
-            self.client.begin_create_or_update,
+            update_client.begin_create_or_update,
             self.context.get_resource_group_name(),
             self.context.get_cluster_name(),
             self.context.get_nodepool_name(),

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

@@ -3100,6 +3100,7 @@ def aks_agentpool_update(
     undrainable_node_behavior=None,
     mode=None,
     scale_down_mode=None,
+    zones=None,
     no_wait=False,
     aks_custom_headers=None,
     allowed_host_ports=None,

src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_agentpool_decorator.py

@@ -50,7 +50,7 @@
     MutuallyExclusiveArgumentError,
     RequiredArgumentMissingError,
 )
-from azure.cli.core.profiles import ResourceType
+from azure.cli.core.profiles import CustomResourceType, ResourceType
 from azure.cli.core.util import get_file_json
 
 class AKSAgentPoolModelsTestCase(unittest.TestCase):
@@ -3570,6 +3570,53 @@ def common_update_vm_properties(self):
         )
         self.assertEqual(dec_agentpool_1, grond_truth_agentpool_1)
 
+    def common_update_zones(self):
+        # --zones auto replaces the current availability zone setting
+        dec_1 = AKSAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"zones": ["auto"]},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        with self.assertRaises(CLIInternalError):
+            dec_1.update_zones(None)
+        agentpool_1 = self.create_initialized_agentpool_instance(availability_zones=None)
+        dec_1.context.attach_agentpool(agentpool_1)
+        dec_agentpool_1 = dec_1.update_zones(agentpool_1)
+        ground_truth_agentpool_1 = self.create_initialized_agentpool_instance(availability_zones=["auto"])
+        self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1)
+
+        # Explicit zone lists are passed through unchanged
+        dec_2 = AKSAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"zones": ["1", "2", "3"]},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        agentpool_2 = self.create_initialized_agentpool_instance(availability_zones=["auto"])
+        dec_2.context.attach_agentpool(agentpool_2)
+        dec_agentpool_2 = dec_2.update_zones(agentpool_2)
+        ground_truth_agentpool_2 = self.create_initialized_agentpool_instance(
+            availability_zones=["1", "2", "3"]
+        )
+        self.assertEqual(dec_agentpool_2, ground_truth_agentpool_2)
+
+        # Omitting --zones preserves the existing setting
+        dec_3 = AKSAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"zones": None},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        agentpool_3 = self.create_initialized_agentpool_instance(availability_zones=["1", "2"])
+        dec_3.context.attach_agentpool(agentpool_3)
+        dec_agentpool_3 = dec_3.update_zones(agentpool_3)
+        ground_truth_agentpool_3 = self.create_initialized_agentpool_instance(availability_zones=["1", "2"])
+        self.assertEqual(dec_agentpool_3, ground_truth_agentpool_3)
+
     def common_update_gpu_profile(self):
         dec_1 = AKSAgentPoolUpdateDecorator(
             self.cmd,
@@ -3892,7 +3939,10 @@ def test_update_upgrade_settings(self):
 
     def test_update_vm_properties(self):
         self.common_update_vm_properties()
-
+
+    def test_update_zones(self):
+        self.common_update_zones()
+
     def test_update_fips_image(self):
         self.common_update_fips_image()
 
@@ -3948,6 +3998,31 @@ def test_update_agentpool_profile_default(self):
         )
         self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1)
 
+        # Explicit zones compose with another nodepool update property in the full update flow
+        raw_param_dict["zones"] = ["auto"]
+        raw_param_dict["scale_down_mode"] = "Deallocate"
+        dec_2 = AKSAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            raw_param_dict,
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        self.client.get = Mock(
+            return_value=self.create_initialized_agentpool_instance(
+                nodepool_name="test_nodepool_name",
+                availability_zones=None,
+                scale_down_mode="Delete",
+            )
+        )
+        dec_agentpool_2 = dec_2.update_agentpool_profile_default()
+        ground_truth_agentpool_2 = self.create_initialized_agentpool_instance(
+            nodepool_name="test_nodepool_name",
+            availability_zones=["auto"],
+            scale_down_mode="Deallocate",
+        )
+        self.assertEqual(dec_agentpool_2, ground_truth_agentpool_2)
+
         dec_1.context.raw_param.print_usage_statistics()
 
     def test_update_agentpool(self):
@@ -3965,10 +4040,97 @@ def test_update_agentpool(self):
         # fail on passing the wrong agentpool object
         with self.assertRaises(CLIInternalError):
             dec_1.update_agentpool(None)
-        agentpool_1 = self.create_initialized_agentpool_instance(nodepool_name="test_nodepool_name")
+        agentpool_1 = self.create_initialized_agentpool_instance(
+            nodepool_name="test_nodepool_name", availability_zones=["auto"]
+        )
         dec_1.context.attach_agentpool(agentpool_1)
-        with patch("azure.cli.command_modules.acs.agentpool_decorator.sdk_no_wait") as put_agentpool:
+        with patch(
+            "azure.cli.command_modules.acs.agentpool_decorator.get_mgmt_service_client"
+        ) as get_client, patch(
+            "azure.cli.command_modules.acs.agentpool_decorator.sdk_no_wait"
+        ) as put_agentpool:
             dec_1.update_agentpool(agentpool_1)
+        get_client.assert_not_called()
+        put_agentpool.assert_called_once_with(
+            False,
+            self.client.begin_create_or_update,
+            "test_resource_group_name",
+            "test_cluster_name",
+            "test_nodepool_name",
+            agentpool_1,
+            headers={},
+        )
+
+    def test_update_agentpool_with_zones_uses_preview_api(self):
+        dec_1 = AKSAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {
+                "resource_group_name": "test_resource_group_name",
+                "cluster_name": "test_cluster_name",
+                "nodepool_name": "test_nodepool_name",
+                "zones": ["auto"],
+            },
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        agentpool_1 = self.create_initialized_agentpool_instance(
+            nodepool_name="test_nodepool_name", availability_zones=["auto"]
+        )
+        dec_1.context.attach_agentpool(agentpool_1)
+        preview_agentpool_client = Mock()
+        preview_client = Mock(agent_pools=preview_agentpool_client)
+
+        with patch(
+            "azure.cli.command_modules.acs.agentpool_decorator.get_mgmt_service_client",
+            return_value=preview_client,
+        ) as get_client, patch(
+            "azure.cli.command_modules.acs.agentpool_decorator.sdk_no_wait"
+        ) as put_agentpool:
+            dec_1.update_agentpool(agentpool_1)
+
+        get_client.assert_called_once_with(
+            self.cmd.cli_ctx,
+            ResourceType.MGMT_CONTAINERSERVICE,
+            api_version="2026-01-02-preview",
+        )
+        put_agentpool.assert_called_once_with(
+            False,
+            preview_agentpool_client.begin_create_or_update,
+            "test_resource_group_name",
+            "test_cluster_name",
+            "test_nodepool_name",
+            agentpool_1,
+            headers={},
+        )
+
+    def test_update_agentpool_with_zones_keeps_custom_resource_client(self):
+        dec_1 = AKSAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {
+                "resource_group_name": "test_resource_group_name",
+                "cluster_name": "test_cluster_name",
+                "nodepool_name": "test_nodepool_name",
+                "zones": ["auto"],
+            },
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        dec_1.resource_type = CustomResourceType("azext_aks_preview.vendored_sdks", "ContainerServiceClient")
+        agentpool_1 = self.create_initialized_agentpool_instance(
+            nodepool_name="test_nodepool_name", availability_zones=["auto"]
+        )
+        dec_1.context.attach_agentpool(agentpool_1)
+
+        with patch(
+            "azure.cli.command_modules.acs.agentpool_decorator.get_mgmt_service_client"
+        ) as get_client, patch(
+            "azure.cli.command_modules.acs.agentpool_decorator.sdk_no_wait"
+        ) as put_agentpool:
+            dec_1.update_agentpool(agentpool_1)
+
+        get_client.assert_not_called()
         put_agentpool.assert_called_once_with(
             False,
             self.client.begin_create_or_update,
@@ -4086,6 +4248,9 @@ def test_update_agentpool_profile_default(self):
     def test_update_vm_properties(self):
         self.common_update_vm_properties()
 
+    def test_update_zones(self):
+        self.common_update_zones()
+
     def test_update_localdns_profile(self):
         self.common_update_localdns_profile()