aks-preview: Support VMSS agent pool VM size resize via nodepool update

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`.
 
 19.0.0b28
 +++++++

src/aks-preview/azext_aks_preview/_help.py

@@ -2478,7 +2478,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".
@@ -2513,6 +2513,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

@@ -2205,6 +2205,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

@@ -1848,6 +1848,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)
@@ -1910,6 +1934,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)
 
@@ -1940,13 +1967,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

@@ -2650,6 +2650,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(
@@ -2990,6 +3053,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()
 
@@ -3086,6 +3152,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()