[AKS] az aks create/update: Add new parameter --enable-ai-toolchain-operator to enable Kaito addon (#31485)

Author: helayoty

linter_exclusions.yml

@@ -285,6 +285,9 @@ aks create:
     azure_keyvault_kms_key_vault_network_access:
       rule_exclusions:
         - option_length_too_long
+    enable_ai_toolchain_operator:
+      rule_exclusions:
+      - option_length_too_long
     azure_keyvault_kms_key_vault_resource_id:
       rule_exclusions:
         - option_length_too_long
@@ -362,6 +365,12 @@ aks update:
     azure_keyvault_kms_key_vault_network_access:
       rule_exclusions:
         - option_length_too_long
+    enable_ai_toolchain_operator:
+      rule_exclusions:
+        - option_length_too_long
+    disable_ai_toolchain_operator:
+      rule_exclusions:
+        - option_length_too_long
     azure_keyvault_kms_key_vault_resource_id:
       rule_exclusions:
         - option_length_too_long

src/azure-cli/azure/cli/command_modules/acs/_help.py

@@ -590,6 +590,9 @@
     short-summary: Configure artifact source when bootstraping the cluster.
     long-summary: |
         The artifacts include the addon image. Use "Direct" to download artifacts from MCR, "Cache" to downalod artifacts from Azure Container Registry.
+  - name: --enable-ai-toolchain-operator
+    type: bool
+    short-summary: Enable AI toolchain operator to the cluster.
   - name: --bootstrap-container-registry-resource-id
     type: string
     short-summary: Configure container registry resource ID. Must use "Cache" as bootstrap artifact source.
@@ -1048,6 +1051,12 @@
     short-summary: Configure artifact source when bootstraping the cluster.
     long-summary: |
         The artifacts include the addon image. Use "Direct" to download artifacts from MCR, "Cache" to downalod artifacts from Azure Container Registry.
+  - name: --enable-ai-toolchain-operator
+    type: bool
+    short-summary: Enable AI toolchain operator to the cluster
+  - name: --disable-ai-toolchain-operator
+    type: bool
+    short-summary: Disable AI toolchain operator.
   - name: --bootstrap-container-registry-resource-id
     type: string
     short-summary: Configure container registry resource ID. Must use "Cache" as bootstrap artifact source.

src/azure-cli/azure/cli/command_modules/acs/_params.py

@@ -434,6 +434,7 @@ def load_arguments(self, _):
         c.argument('rotation_poll_interval')
         c.argument('enable_sgxquotehelper', action='store_true')
         c.argument('enable_app_routing', action="store_true")
+        c.argument('enable_ai_toolchain_operator', action='store_true')
         c.argument(
             "app_routing_default_nginx_controller",
             arg_type=get_enum_type(app_routing_nginx_configs),
@@ -633,6 +634,8 @@ def load_arguments(self, _):
         # addons
         c.argument('enable_secret_rotation', action='store_true')
         c.argument('disable_secret_rotation', action='store_true', validator=validate_keyvault_secrets_provider_disable_and_enable_parameters)
+        c.argument('enable_ai_toolchain_operator', action='store_true')
+        c.argument('disable_ai_toolchain_operator', action='store_true')
         c.argument('rotation_poll_interval')
         c.argument('enable_static_egress_gateway', action='store_true')
         c.argument('disable_static_egress_gateway', action='store_true')

src/azure-cli/azure/cli/command_modules/acs/custom.py

@@ -602,6 +602,7 @@ def aks_create(
     enable_secret_rotation=False,
     rotation_poll_interval=None,
     enable_app_routing=False,
+    enable_ai_toolchain_operator=False,
     app_routing_default_nginx_controller=None,
     enable_static_egress_gateway=False,
     # nodepool paramerters
@@ -803,6 +804,8 @@ def aks_update(
     # addons
     enable_secret_rotation=False,
     disable_secret_rotation=False,
+    enable_ai_toolchain_operator=False,
+    disable_ai_toolchain_operator=False,
     rotation_poll_interval=None,
     enable_static_egress_gateway=False,
     disable_static_egress_gateway=False,

src/azure-cli/azure/cli/command_modules/acs/managed_cluster_decorator.py

@@ -5414,6 +5414,32 @@ def get_upgrade_override_until(self) -> Union[str, None]:
         # this parameter does not need validation
         return self.raw_param.get("upgrade_override_until")
 
+    def get_ai_toolchain_operator(self, enable_validation: bool = False) -> bool:
+        """Internal function to obtain the value of enable_ai_toolchain_operator.
+        When enabled, if both enable_ai_toolchain_operator and
+        disable_ai_toolchain_operator are specified, raise
+        a MutuallyExclusiveArgumentError.
+        :return: bool
+        """
+        enable_ai_toolchain_operator = self.raw_param.get("enable_ai_toolchain_operator")
+        # This parameter does not need dynamic completion.
+        if enable_validation:
+            if enable_ai_toolchain_operator and self.get_disable_ai_toolchain_operator():
+                raise MutuallyExclusiveArgumentError(
+                    "Cannot specify --enable-ai-toolchain-operator and "
+                    "--disable-ai-toolchain-operator at the same time. "
+                )
+
+        return enable_ai_toolchain_operator
+
+    def get_disable_ai_toolchain_operator(self) -> bool:
+        """Obtain the value of disable_ai_toolchain_operator.
+        :return: bool
+        """
+        # Note: No need to check for mutually exclusive parameter with enable-ai-toolchain-operator here
+        # because it's already checked in get_ai_toolchain_operator
+        return self.raw_param.get("disable_ai_toolchain_operator")
+
     def _get_enable_cost_analysis(self, enable_validation: bool = False) -> bool:
         """Internal function to obtain the value of enable_cost_analysis.
         When enabled, if both enable_cost_analysis and disable_cost_analysis are
@@ -6772,6 +6798,18 @@ def set_up_ingress_web_app_routing(self, mc: ManagedCluster) -> ManagedCluster:
 
         return mc
 
+    def set_up_ai_toolchain_operator(self, mc: ManagedCluster) -> ManagedCluster:
+        self._ensure_mc(mc)
+
+        if self.context.get_ai_toolchain_operator(enable_validation=True):
+            if mc.ai_toolchain_operator_profile is None:
+                mc.ai_toolchain_operator_profile = self.models.ManagedClusterAIToolchainOperatorProfile()  # pylint: disable=no-member
+            # set enabled
+            mc.ai_toolchain_operator_profile.enabled = True
+
+        # Default is disabled so no need to worry about that here
+        return mc
+
     def set_up_cost_analysis(self, mc: ManagedCluster) -> ManagedCluster:
         self._ensure_mc(mc)
 
@@ -6927,6 +6965,8 @@ def construct_mc_profile_default(self, bypass_restore_defaults: bool = False) ->
         mc = self.set_up_metrics_profile(mc)
         # set up node resource group profile
         mc = self.set_up_node_resource_group_profile(mc)
+        # set up AI toolchain operator
+        mc = self.set_up_ai_toolchain_operator(mc)
         # set up bootstrap profile
         mc = self.set_up_bootstrap_profile(mc)
         # set up static egress gateway profile
@@ -8709,6 +8749,23 @@ def update_azure_container_storage(self, mc: ManagedCluster) -> ManagedCluster:
 
         return mc
 
+    def update_ai_toolchain_operator(self, mc: ManagedCluster) -> ManagedCluster:
+        """Updates the aiToolchainOperatorProfile field of the managed cluster
+        :return: the ManagedCluster object
+        """
+
+        if self.context.get_ai_toolchain_operator(enable_validation=True):
+            if mc.ai_toolchain_operator_profile is None:
+                mc.ai_toolchain_operator_profile = self.models.ManagedClusterAIToolchainOperatorProfile()  # pylint: disable=no-member
+            mc.ai_toolchain_operator_profile.enabled = True
+
+        if self.context.get_disable_ai_toolchain_operator():
+            if mc.ai_toolchain_operator_profile is None:
+                mc.ai_toolchain_operator_profile = self.models.ManagedClusterAIToolchainOperatorProfile()  # pylint: disable=no-member
+            mc.ai_toolchain_operator_profile.enabled = False
+
+        return mc
+
     def update_cost_analysis(self, mc: ManagedCluster) -> ManagedCluster:
         self._ensure_mc(mc)
 
@@ -8898,6 +8955,8 @@ def update_mc_profile_default(self) -> ManagedCluster:
         mc = self.update_metrics_profile(mc)
         # update node resource group profile
         mc = self.update_node_resource_group_profile(mc)
+        # update AI toolchain operator
+        mc = self.update_ai_toolchain_operator(mc)
         # update bootstrap profile
         mc = self.update_bootstrap_profile(mc)
         # update static egress gateway