[AKS] Add Node Initialization Taints preview CLI (#7452)

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
 +++++++
+* Add `--nodepool-initialization-taints` to `az aks create` and `az aks update`.
 
 3.0.0b2
 +++++++

src/aks-preview/azcli_aks_live_test/configs/ext_matrix_default.json

@@ -51,6 +51,9 @@
             "test_aks_create_update_secure_boot_flow",
             "test_aks_create_update_vtpm_flow"
         ],
+        "nodepool permanent and initialization taints, missing feature registration": [
+            "test_aks_create_cluster_with_taints"
+        ],
         "automatic, missing feature registration & toggle": [
             "test_aks_automatic_sku"
         ]

src/aks-preview/azext_aks_preview/_help.py

@@ -587,6 +587,9 @@
         - name: --nodepool-taints
           type: string
           short-summary: The node taints for all node pools in this cluster.
+        - name: --node-init-taints --nodepool-initialization-taints
+          type: string
+          short-summary: The node initialization taints for node pools created with aks create operation.
         - name: --enable-cost-analysis
           type: bool
           short-summary: Enable exporting Kubernetes Namespace and Deployment details to the Cost Analysis views in the Azure portal. For more information see aka.ms/aks/docs/cost-analysis.
@@ -1179,6 +1182,9 @@
         - name: --ssh-access
           type: string
           short-summary: Update SSH setting for all node pools in this cluster. Use "disabled" to disable SSH access, "localuser" to enable SSH access using private key.
+        - name: --node-init-taints --nodepool-initialization-taints
+          type: string
+          short-summary: The node initialization taints for all node pools in cluster.
     examples:
       - name: Reconcile the cluster back to its current state.
         text: az aks update -g MyResourceGroup -n MyManagedCluster

src/aks-preview/azext_aks_preview/_params.py

@@ -647,6 +647,16 @@ def load_arguments(self, _):
             ),
         )
         c.argument("nodepool_taints", validator=validate_nodepool_taints)
+        c.argument(
+            "nodepool_initialization_taints",
+            options_list=["--nodepool-initialization-taints", "--node-init-taints"],
+            is_preview=True,
+            validator=validate_nodepool_taints,
+            help=(
+                "Comma-separated taints: <key1>=<value1>:<effect1>,<key2>=<value2>:<effect2>. "
+                "Pass \"\" to clear existing taints."
+            ),
+        )
         c.argument("node_osdisk_type", arg_type=get_enum_type(node_os_disk_types))
         c.argument("node_osdisk_size", type=int)
         c.argument("max_pods", type=int, options_list=["--max-pods", "-m"])
@@ -1065,6 +1075,16 @@ def load_arguments(self, _):
             )
         )
         c.argument("nodepool_taints", validator=validate_nodepool_taints)
+        c.argument(
+            "nodepool_initialization_taints",
+            options_list=["--nodepool-initialization-taints", "--node-init-taints"],
+            is_preview=True,
+            validator=validate_nodepool_taints,
+            help=(
+                "Comma-separated taints: <key1>=<value1>:<effect1>,<key2>=<value2>:<effect2>. "
+                "Pass \"\" to clear existing taints."
+            ),
+        )
         # misc
         c.argument(
             "yes",

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -382,6 +382,28 @@ def get_node_taints(self) -> Union[List[str], None]:
         # this parameter does not need validation
         return node_taints
 
+    def get_node_initialization_taints(self) -> Union[List[str], None]:
+        """Obtain the value of node_initialization_taints.
+
+        :return: empty list, list of strings or None
+        """
+        # read the original value passed by the command
+        node_init_taints = self.raw_param.get("nodepool_initialization_taints")
+        # normalize, default is an empty list
+        if node_init_taints is not None:
+            node_init_taints = [x.strip() for x in (node_init_taints.split(",") if node_init_taints else [""])]
+        # keep None as None for update mode
+        if node_init_taints is None and self.decorator_mode == DecoratorMode.CREATE:
+            node_init_taints = []
+
+        # In create mode, try to read the property value corresponding to the parameter from the `agentpool` object
+        if self.decorator_mode == DecoratorMode.CREATE:
+            if self.agentpool and self.agentpool.node_initialization_taints is not None:
+                node_init_taints = self.agentpool.node_initialization_taints
+
+        # this parameter does not need validation
+        return node_init_taints
+
     def get_drain_timeout(self):
         """Obtain the value of drain_timeout.
 
@@ -707,6 +729,15 @@ def set_up_taints(self, agentpool: AgentPool) -> AgentPool:
         agentpool.node_taints = self.context.get_node_taints()
         return agentpool
 
+    def set_up_init_taints(self, agentpool: AgentPool) -> AgentPool:
+        """Set up label, tag, taint for the AgentPool object.
+
+        :return: the AgentPool object
+        """
+        self._ensure_agentpool(agentpool)
+        agentpool.node_initialization_taints = self.context.get_node_initialization_taints()
+        return agentpool
+
     def set_up_artifact_streaming(self, agentpool: AgentPool) -> AgentPool:
         """Set up artifact streaming property for the AgentPool object."""
         self._ensure_agentpool(agentpool)
@@ -797,6 +828,8 @@ def construct_agentpool_profile_preview(self) -> AgentPool:
         agentpool = self.set_up_agentpool_network_profile(agentpool)
         # set up taints
         agentpool = self.set_up_taints(agentpool)
+        # set up initialization taints
+        agentpool = self.set_up_init_taints(agentpool)
         # set up artifact streaming
         agentpool = self.set_up_artifact_streaming(agentpool)
         # set up skip_gpu_driver_install

src/aks-preview/azext_aks_preview/custom.py

@@ -557,6 +557,7 @@ def aks_create(
     nodepool_tags=None,
     nodepool_labels=None,
     nodepool_taints=None,
+    nodepool_initialization_taints=None,
     node_osdisk_type=None,
     node_osdisk_size=0,
     vm_set_type=None,
@@ -756,6 +757,7 @@ def aks_update(
     max_count=None,
     nodepool_labels=None,
     nodepool_taints=None,
+    nodepool_initialization_taints=None,
     # misc
     yes=False,
     no_wait=False,

src/aks-preview/azext_aks_preview/managed_cluster_decorator.py

@@ -2531,6 +2531,13 @@ def get_nodepool_taints(self) -> Union[List[str], None]:
         """
         return self.agentpool_context.get_node_taints()
 
+    def get_nodepool_initialization_taints(self) -> Union[List[str], None]:
+        """Obtain the value of nodepool_initialization_taints.
+
+        :return: dictionary or None
+        """
+        return self.agentpool_context.get_node_initialization_taints()
+
     def _get_enable_cost_analysis(self, enable_validation: bool = False) -> bool:
         """Internal function to obtain the value of enable_cost_analysis.
 
@@ -3610,6 +3617,7 @@ def get_special_parameter_default_value_pairs_list(self) -> List[Tuple[Any, Any]
             (self.context.get_api_server_authorized_ip_ranges(), None),
             (self.context.get_nodepool_labels(), None),
             (self.context.get_nodepool_taints(), None),
+            (self.context.get_nodepool_initialization_taints(), None),
             (self.context.raw_param.get("upgrade_settings"), None),
             (self.context.get_load_balancer_managed_outbound_ip_count(), None),
             (self.context.get_load_balancer_managed_outbound_ipv6_count(), None),
@@ -4543,6 +4551,22 @@ def update_nodepool_taints_mc(self, mc: ManagedCluster) -> ManagedCluster:
                 agent_profile.node_taints = nodepool_taints
         return mc
 
+    def update_nodepool_initialization_taints_mc(self, mc: ManagedCluster) -> ManagedCluster:
+        self._ensure_mc(mc)
+
+        if not mc.agent_pool_profiles:
+            raise UnknownError(
+                "Encounter an unexpected error while getting agent pool profiles from the cluster in the process of "
+                "updating agentpool profile."
+            )
+
+        # update nodepool taints for all nodepools
+        nodepool_initialization_taints = self.context.get_nodepool_initialization_taints()
+        if nodepool_initialization_taints is not None:
+            for agent_profile in mc.agent_pool_profiles:
+                agent_profile.node_initialization_taints = nodepool_initialization_taints
+        return mc
+
     def update_cost_analysis(self, mc: ManagedCluster) -> ManagedCluster:
         self._ensure_mc(mc)
 
@@ -4830,6 +4854,8 @@ def update_mc_profile_preview(self) -> ManagedCluster:
         mc = self.update_upgrade_settings(mc)
         # update nodepool taints
         mc = self.update_nodepool_taints_mc(mc)
+        # update nodepool initialization taints
+        mc = self.update_nodepool_initialization_taints_mc(mc)
         # update network_observability in network_profile
         mc = self.update_enable_network_observability_in_network_profile(mc)
         # update kubernetes support plan