aks-preview: Support VMSS agent pool VM size resize via nodepool update (#9732)

Author: wenhug

src/aks-preview/HISTORY.rst

@@ -11,6 +11,7 @@ To release a new version, please select a new version number (usually plus 1 to
 
 Pending
 +++++++
+* `az aks nodepool update`: Support `--node-vm-size` to resize VM size of an existing VMSS-based agent pool (preview). Requires AFEC registration `Microsoft.ContainerService/AgentPoolVMSSResize`.
 
 20.0.0b3
 ++++++

src/aks-preview/azext_aks_preview/_help.py

@@ -2484,7 +2484,7 @@
           short-summary: Set the localDNS Profile for a nodepool with a JSON config file.
         - name: --node-vm-size -s
           type: string
-          short-summary: VM size for Kubernetes nodes. Only configurable when updating the autoscale settings of a VirtualMachines node pool.
+          short-summary: VM size for Kubernetes nodes. For VMSS pools, changing this triggers a rolling upgrade to replace nodes with the new size (preview). For VirtualMachines pools, only configurable when updating autoscale settings.
         - name: --upgrade-strategy
           type: string
           short-summary: Upgrade strategy for the node pool. Allowed values are "Rolling" or "BlueGreen". Default is "Rolling".
@@ -2519,6 +2519,8 @@
         text: az aks nodepool update --mode System -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster
       - name: Update cluster autoscaler vm size, min-count and max-count for virtual machines node pool
         text: az aks nodepool update -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster --update-cluster-autoscaler --node-vm-size "Standard_D2s_v3" --min-count 2 --max-count 4
+      - name: Resize VM size for a VMSS node pool (preview, requires AFEC registration)
+        text: az aks nodepool update -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster --node-vm-size Standard_D4s_v3
       - name: Update a node pool with blue-green upgrade settings
         text: az aks nodepool update -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster --drain-batch-size 50% --drain-timeout-bg 5 --batch-soak-duration 10 --final-soak-duration 10
 """

src/aks-preview/azext_aks_preview/_params.py

@@ -2299,6 +2299,7 @@ def load_arguments(self, _):
             "node_vm_size",
             options_list=["--node-vm-size", "-s"],
             completer=get_vm_size_completion_list,
+            is_preview=True,
         )
         c.argument(
             "gpu_driver",

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -1944,6 +1944,30 @@ def update_fips_image(self, agentpool: AgentPool) -> AgentPool:
 
         return agentpool
 
+    def update_vm_size(self, agentpool: AgentPool) -> AgentPool:
+        """Update VM size for the AgentPool object.
+
+        Allows changing the VM size (SKU) of an existing VMSS-based agent pool.
+        The RP will perform a rolling upgrade (surge new nodes, drain old, delete old)
+        to replace nodes with the new VM size.
+
+        Note: This is only for VMSS pools. VMs pools handle VM size changes through
+        the autoscaler update path (update_auto_scaler_properties_vms).
+
+        :return: the AgentPool object
+        """
+        self._ensure_agentpool(agentpool)
+
+        # Skip for VirtualMachines pools - they handle VM size via autoscaler path
+        if self.context.get_vm_set_type() == CONST_VIRTUAL_MACHINES:
+            return agentpool
+
+        node_vm_size = self.context.raw_param.get("node_vm_size")
+        if node_vm_size:
+            agentpool.vm_size = node_vm_size
+
+        return agentpool
+
     def update_localdns_profile(self, agentpool: AgentPool) -> AgentPool:
         """Update local DNS profile for the AgentPool object if provided via --localdns-config."""
         self._ensure_agentpool(agentpool)
@@ -2006,6 +2030,9 @@ def update_agentpool_profile_preview(self, agentpools: List[AgentPool] = None) -
         # update ssh access
         agentpool = self.update_ssh_access(agentpool)
 
+        # update vm size for VMSS pools
+        agentpool = self.update_vm_size(agentpool)
+
         # update local DNS profile
         agentpool = self.update_localdns_profile(agentpool)
 
@@ -2039,13 +2066,6 @@ def update_auto_scaler_properties(self, agentpool: AgentPool) -> AgentPool:
         if self.context.get_vm_set_type() == CONST_VIRTUAL_MACHINES:
             return agentpool
 
-        vm_size = self.context.raw_param.get("node_vm_size")
-        if vm_size is not None:
-            raise InvalidArgumentValueError(
-                "--node-vm-size can only be used with virtual machines agentpools. "
-                "Updating VM size is not supported for virtual machine scale set agentpools."
-            )
-
         (
             update_cluster_autoscaler,
             enable_cluster_autoscaler,

src/aks-preview/azext_aks_preview/tests/latest/test_agentpool_decorator.py

@@ -2773,6 +2773,69 @@ def common_update_fips_image(self):
         with self.assertRaises(MutuallyExclusiveArgumentError):
             dec_3.update_fips_image(agentpool_2)
 
+    def common_update_vm_size(self):
+        # Test case 1: No node_vm_size provided (should not change agentpool)
+        dec_1 = AKSPreviewAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"node_vm_size": None},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        # fail on passing the wrong agentpool object
+        with self.assertRaises(CLIInternalError):
+            dec_1.update_vm_size(None)
+
+        agentpool_1 = self.create_initialized_agentpool_instance(
+            vm_size="Standard_D2s_v3"
+        )
+        dec_1.context.attach_agentpool(agentpool_1)
+        dec_agentpool_1 = dec_1.update_vm_size(agentpool_1)
+        ground_truth_agentpool_1 = self.create_initialized_agentpool_instance(
+            vm_size="Standard_D2s_v3"
+        )
+        self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1)
+
+        # Test case 2: node_vm_size provided (should update agentpool)
+        dec_2 = AKSPreviewAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"node_vm_size": "Standard_D4s_v3"},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        agentpool_2 = self.create_initialized_agentpool_instance(
+            vm_size="Standard_D2s_v3"
+        )
+        dec_2.context.attach_agentpool(agentpool_2)
+        dec_agentpool_2 = dec_2.update_vm_size(agentpool_2)
+        ground_truth_agentpool_2 = self.create_initialized_agentpool_instance(
+            vm_size="Standard_D4s_v3"
+        )
+        self.assertEqual(dec_agentpool_2, ground_truth_agentpool_2)
+
+        # Test case 3: VirtualMachines pool with node_vm_size provided (should be no-op)
+        dec_3 = AKSPreviewAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"node_vm_size": "Standard_D4s_v3"},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        agentpool_3 = self.create_initialized_agentpool_instance(
+            vm_size="Standard_D2s_v3"
+        )
+        # Set pool type to VirtualMachines - use the correct attribute based on decorator mode
+        from azure.cli.command_modules.acs._consts import AgentPoolDecoratorMode
+        if self.agentpool_decorator_mode == AgentPoolDecoratorMode.MANAGED_CLUSTER:
+            agentpool_3.type = CONST_VIRTUAL_MACHINES
+        else:
+            agentpool_3.type_properties_type = CONST_VIRTUAL_MACHINES
+        dec_3.context.attach_agentpool(agentpool_3)
+        dec_agentpool_3 = dec_3.update_vm_size(agentpool_3)
+        # vm_size should remain unchanged for VMs pools
+        self.assertEqual(dec_agentpool_3.vm_size, "Standard_D2s_v3")
+
     def common_update_upgrade_strategy(self):
         # Test case 1: No upgrade strategy provided (should not change agentpool)
         dec_1 = AKSPreviewAgentPoolUpdateDecorator(
@@ -3113,6 +3176,9 @@ def test_update_vtpm(self):
     def test_update_fips_image(self):
         self.common_update_fips_image()
 
+    def test_update_vm_size(self):
+        self.common_update_vm_size()
+
     def test_update_upgrade_strategy(self):
         self.common_update_upgrade_strategy()
 
@@ -3209,6 +3275,9 @@ def test_update_vtpm(self):
     def test_update_fips_image(self):
         self.common_update_fips_image()
 
+    def test_update_vm_size(self):
+        self.common_update_vm_size()
+
     def test_update_upgrade_strategy(self):
         self.common_update_upgrade_strategy()
 

src/aks-preview/azext_aks_preview/tests/latest/test_aks_commands.py

@@ -23405,3 +23405,51 @@ def test_aks_list_vm_skus(self):
                 (sku.get("locationInfo") or [{}])[0].get("zones") or []
             )
             assert len(zones) > 0, f"SKU '{sku['name']}' has no zones despite --zone filter"
+
+    @live_only()
+    @AllowLargeResponse()
+    @AKSCustomResourceGroupPreparer(
+        random_name_length=17, name_prefix="clitest", location="centraluseuap"
+    )
+    def test_aks_nodepool_update_vmss_vm_size_resize(
+        self, resource_group, resource_group_location
+    ):
+        """Test VMSS agent pool VM size resize via nodepool update (preview)."""
+        aks_name = self.create_random_name("cliakstest", 16)
+        nodepool_name = "nodepool1"
+        self.kwargs.update(
+            {
+                "resource_group": resource_group,
+                "name": aks_name,
+                "nodepool_name": nodepool_name,
+                "ssh_key_value": self.generate_ssh_keys(),
+            }
+        )
+
+        # Create cluster with Standard_D2s_v3
+        create_cmd = (
+            "aks create --resource-group={resource_group} --name={name} "
+            "--node-count=1 --node-vm-size Standard_D2s_v3 "
+            "--ssh-key-value={ssh_key_value}"
+        )
+        self.cmd(
+            create_cmd,
+            checks=[
+                self.check("provisioningState", "Succeeded"),
+                self.check("agentPoolProfiles[0].vmSize", "Standard_D2s_v3"),
+            ],
+        )
+
+        # Resize nodepool VM size to Standard_D4s_v3
+        update_cmd = (
+            "aks nodepool update --resource-group={resource_group} "
+            "--cluster-name={name} -n {nodepool_name} "
+            "--node-vm-size Standard_D4s_v3"
+        )
+        self.cmd(
+            update_cmd,
+            checks=[
+                self.check("provisioningState", "Succeeded"),
+                self.check("vmSize", "Standard_D4s_v3"),
+            ],
+        )

src/aks-preview/azext_aks_preview/tests/latest/test_update_agentpool_profile_preview.py

@@ -129,6 +129,7 @@ def test_update_agentpool_profile_preview_default_behavior(self):
         decorator.update_os_sku = Mock(return_value=agentpool)
         decorator.update_fips_image = Mock(return_value=agentpool)
         decorator.update_ssh_access = Mock(return_value=agentpool)
+        decorator.update_vm_size = Mock(return_value=agentpool)
         decorator.update_localdns_profile = Mock(return_value=agentpool)
         decorator.update_auto_scaler_properties_vms = Mock(return_value=agentpool)
         decorator.update_upgrade_strategy = Mock(return_value=agentpool)
@@ -154,6 +155,7 @@ def test_update_agentpool_profile_preview_default_behavior(self):
         decorator.update_os_sku.assert_called_once_with(agentpool)
         decorator.update_fips_image.assert_called_once_with(agentpool)
         decorator.update_ssh_access.assert_called_once_with(agentpool)
+        decorator.update_vm_size.assert_called_once_with(agentpool)
         decorator.update_localdns_profile.assert_called_once_with(agentpool)
         decorator.update_auto_scaler_properties_vms.assert_called_once_with(agentpool)
         decorator.update_upgrade_strategy.assert_called_once_with(agentpool)
@@ -197,6 +199,7 @@ def test_update_agentpool_profile_preview_with_agentpools_parameter(self):
         decorator.update_os_sku = Mock(return_value=agentpool)
         decorator.update_fips_image = Mock(return_value=agentpool)
         decorator.update_ssh_access = Mock(return_value=agentpool)
+        decorator.update_vm_size = Mock(return_value=agentpool)
         decorator.update_localdns_profile = Mock(return_value=agentpool)
         decorator.update_auto_scaler_properties_vms = Mock(return_value=agentpool)
         decorator.update_upgrade_strategy = Mock(return_value=agentpool)
@@ -359,6 +362,7 @@ def test_update_agentpool_profile_preview_system_mode_regular_flow(self):
         decorator.update_os_sku = Mock(return_value=agentpool)
         decorator.update_fips_image = Mock(return_value=agentpool)
         decorator.update_ssh_access = Mock(return_value=agentpool)
+        decorator.update_vm_size = Mock(return_value=agentpool)
         decorator.update_localdns_profile = Mock(return_value=agentpool)
         decorator.update_auto_scaler_properties_vms = Mock(return_value=agentpool)
         decorator.update_upgrade_strategy = Mock(return_value=agentpool)
@@ -382,6 +386,7 @@ def test_update_agentpool_profile_preview_system_mode_regular_flow(self):
         decorator.update_os_sku.assert_called_once_with(agentpool)
         decorator.update_fips_image.assert_called_once_with(agentpool)
         decorator.update_ssh_access.assert_called_once_with(agentpool)
+        decorator.update_vm_size.assert_called_once_with(agentpool)
         decorator.update_localdns_profile.assert_called_once_with(agentpool)
         decorator.update_auto_scaler_properties_vms.assert_called_once_with(agentpool)
         decorator.update_upgrade_strategy.assert_called_once_with(agentpool)
@@ -430,6 +435,7 @@ def mock_method(pool):
         decorator.update_os_sku = create_mock_update_method("update_os_sku")
         decorator.update_fips_image = create_mock_update_method("update_fips_image")
         decorator.update_ssh_access = create_mock_update_method("update_ssh_access")
+        decorator.update_vm_size = create_mock_update_method("update_vm_size")
         decorator.update_localdns_profile = create_mock_update_method("update_localdns_profile")
         decorator.update_auto_scaler_properties_vms = create_mock_update_method("update_auto_scaler_properties_vms")
         decorator.update_upgrade_strategy = create_mock_update_method("update_upgrade_strategy")
@@ -450,6 +456,7 @@ def mock_method(pool):
             "update_os_sku",
             "update_fips_image",
             "update_ssh_access",
+            "update_vm_size",
             "update_localdns_profile",
             "update_auto_scaler_properties_vms",
             "update_upgrade_strategy",
@@ -500,6 +507,7 @@ def track_and_return(pool):
         decorator.update_os_sku = create_tracking_mock("update_os_sku")
         decorator.update_fips_image = create_tracking_mock("update_fips_image")
         decorator.update_ssh_access = create_tracking_mock("update_ssh_access")
+        decorator.update_vm_size = create_tracking_mock("update_vm_size")
         decorator.update_localdns_profile = create_tracking_mock("update_localdns_profile")
         decorator.update_auto_scaler_properties_vms = create_tracking_mock("update_auto_scaler_properties_vms")
         decorator.update_upgrade_strategy = create_tracking_mock("update_upgrade_strategy")
@@ -567,7 +575,7 @@ def test_update_agentpool_profile_preview_mixed_modes_scenario(self):
                 update_methods = [
                     'update_network_profile', 'update_artifact_streaming', 'update_managed_gpu',
                     'update_secure_boot', 'update_vtpm', 'update_os_sku', 'update_fips_image',
-                    'update_ssh_access', 'update_localdns_profile', 'update_auto_scaler_properties_vms', 
+                    'update_ssh_access', 'update_vm_size', 'update_localdns_profile', 'update_auto_scaler_properties_vms', 
                     'update_upgrade_strategy', 'update_blue_green_upgrade_settings', 'update_gpu_profile',
                     'update_gpu_mig_strategy'
                 ]
@@ -638,6 +646,7 @@ def test_update_agentpool_profile_preview_managed_cluster_mode(self):
         decorator.update_os_sku = Mock(return_value=agentpool)
         decorator.update_fips_image = Mock(return_value=agentpool)
         decorator.update_ssh_access = Mock(return_value=agentpool)
+        decorator.update_vm_size = Mock(return_value=agentpool)
         decorator.update_localdns_profile = Mock(return_value=agentpool)
         decorator.update_auto_scaler_properties_vms = Mock(return_value=agentpool)
         decorator.update_upgrade_strategy = Mock(return_value=agentpool)