aks update node disruption profile flag (#9893)

Author: erinborders

src/aks-preview/HISTORY.rst

@@ -13,6 +13,7 @@ Pending
 +++++++
 * `az aks update`: Add `--control-plane-scaling-size` parameter to update the control plane scaling size on an existing cluster with available sizes 'H2', 'H4', and 'H8'.
 * `az aks bastion`: Fix `--subscription` not being passed to internal `az network bastion tunnel` and bastion discovery commands.
+* `az aks update`: Add `--node-disruption-policy` (preview) to update the node disruption policy for a cluster. Requires AFEC registration `Microsoft.ContainerService/NodeDisruptionProfile`. This is a cluster-level property that applies to all node pools in the cluster.
 
 21.0.0b3
 +++++++

src/aks-preview/azext_aks_preview/_consts.py

@@ -414,3 +414,8 @@
 CONST_K8S_EXTENSION_NAME = "k8s-extension"
 CONST_K8S_EXTENSION_ACTION_MOD_NAME = "azext_k8s_extension.action"
 CONST_K8S_EXTENSION_FORMAT_MOD_NAME = "azext_k8s_extension._format"
+
+# Node Disruption Policy Consts
+CONST_NODE_DISRUPTION_POLICY_ALLOW = "Allow"
+CONST_NODE_DISRUPTION_POLICY_BLOCK = "Block"
+CONST_NODE_DISRUPTION_POLICY_ALLOW_DURING_MAINTENANCE_WINDOW = "AllowDuringMaintenanceWindow"

src/aks-preview/azext_aks_preview/_params.py

@@ -78,6 +78,9 @@
     CONST_NETWORK_PLUGIN_NONE,
     CONST_NETWORK_POD_IP_ALLOCATION_MODE_DYNAMIC_INDIVIDUAL,
     CONST_NETWORK_POD_IP_ALLOCATION_MODE_STATIC_BLOCK,
+    CONST_NODE_DISRUPTION_POLICY_ALLOW,
+    CONST_NODE_DISRUPTION_POLICY_BLOCK,
+    CONST_NODE_DISRUPTION_POLICY_ALLOW_DURING_MAINTENANCE_WINDOW,
     CONST_NODE_IMAGE_UPGRADE_CHANNEL,
     CONST_NODE_OS_CHANNEL_NODE_IMAGE,
     CONST_NODE_OS_CHANNEL_NONE,
@@ -562,6 +565,12 @@
     CONST_UPGRADE_STRATEGY_BLUE_GREEN,
 ]
 
+node_disruption_policies = [
+    CONST_NODE_DISRUPTION_POLICY_ALLOW,
+    CONST_NODE_DISRUPTION_POLICY_BLOCK,
+    CONST_NODE_DISRUPTION_POLICY_ALLOW_DURING_MAINTENANCE_WINDOW,
+]
+
 
 def load_arguments(self, _):
     acr_arg_type = CLIArgumentType(metavar="ACR_NAME_OR_RESOURCE_ID")
@@ -1939,6 +1948,12 @@ def load_arguments(self, _):
                  "Available values are 'H2', 'H4', and 'H8'. "
                  "The control plane scaling profile must already be enabled on the cluster.",
         )
+        c.argument(
+            "node_disruption_policy",
+            arg_type=get_enum_type(node_disruption_policies),
+            is_preview=True,
+            help="Set the node disruption policy for the cluster.",
+        )
 
     with self.argument_context("aks delete") as c:
         c.argument("if_match")

src/aks-preview/azext_aks_preview/custom.py

@@ -1439,6 +1439,8 @@ def aks_update(
     # health monitor
     enable_continuous_control_plane_and_addon_monitor=False,
     disable_continuous_control_plane_and_addon_monitor=False,
+    # node disruption policy
+    node_disruption_policy=None,
     # control plane scaling
     control_plane_scaling_size=None,
 ):

src/aks-preview/azext_aks_preview/managed_cluster_decorator.py

@@ -3818,6 +3818,11 @@ def get_node_provisioning_mode(self) -> Union[str, None]:
         """
         return self.raw_param.get("node_provisioning_mode")
 
+    def get_node_disruption_policy(self) -> Union[str, None]:
+        """Obtain the value of node_disruption_policy.
+        """
+        return self.raw_param.get("node_disruption_policy")
+
     def get_node_provisioning_default_pools(self) -> Union[str, None]:
         """Obtain the value of node_provisioning_default_pools.
         """
@@ -7722,6 +7727,25 @@ def update_node_provisioning_profile(self, mc: ManagedCluster) -> ManagedCluster
 
         return mc
 
+    def update_node_disruption_policy(self, mc: ManagedCluster) -> ManagedCluster:
+        """Updates the nodeDisruptionPolicy field of the managed cluster
+
+        :return: the ManagedCluster object
+        """
+        self._ensure_mc(mc)
+
+        policy = self.context.get_node_disruption_policy()
+        if policy is not None:
+            if mc.node_disruption_profile is None:
+                mc.node_disruption_profile = (
+                    self.models.NodeDisruptionProfile()  # pylint: disable=no-member
+                )
+
+            # set policy
+            mc.node_disruption_profile.node_disruption_policy = policy
+
+        return mc
+
     def update_ai_toolchain_operator(self, mc: ManagedCluster) -> ManagedCluster:
         """Updates the aiToolchainOperatorProfile field of the managed cluster
 
@@ -8281,6 +8305,8 @@ def update_mc_profile_preview(self) -> ManagedCluster:
         mc = self.update_http_proxy_enabled(mc)
         # update user-defined scheduler configuration for kube-scheduler upstream
         mc = self.update_upstream_kubescheduler_user_configuration(mc)
+        # update node disruption policy
+        mc = self.update_node_disruption_policy(mc)
         # update control plane scaling profile
         mc = self.update_control_plane_scaling_profile(mc)
         # update ManagedSystem pools, must at end