diff --git a/src/azure-cli/azure/cli/command_modules/acs/_consts.py b/src/azure-cli/azure/cli/command_modules/acs/_consts.py index 227790750ee..c71c362b8a4 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/_consts.py +++ b/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" diff --git a/src/azure-cli/azure/cli/command_modules/acs/_help.py b/src/azure-cli/azure/cli/command_modules/acs/_help.py index e0b7ae94fcd..45b752ad829 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/_help.py +++ b/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"] = """ diff --git a/src/azure-cli/azure/cli/command_modules/acs/_params.py b/src/azure-cli/azure/cli/command_modules/acs/_params.py index 09056477334..a031d17ecc0 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/_params.py +++ b/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) diff --git a/src/azure-cli/azure/cli/command_modules/acs/agentpool_decorator.py b/src/azure-cli/azure/cli/command_modules/acs/agentpool_decorator.py index 70be49f2106..756f25dea41 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/agentpool_decorator.py +++ b/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(), diff --git a/src/azure-cli/azure/cli/command_modules/acs/custom.py b/src/azure-cli/azure/cli/command_modules/acs/custom.py index a0008c954d2..b3f0ce9b747 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/custom.py +++ b/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, diff --git a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_agentpool_decorator.py b/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_agentpool_decorator.py index f4b98145f5f..891cd9940c3 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_agentpool_decorator.py +++ b/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()