[AKS] az command aks: Add KataVmIsolation option for --workload-runtime parameter (#32020)

Author: Bickor

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

@@ -261,6 +261,9 @@
 CONST_NODE_PROVISIONING_DEFAULT_POOLS_NONE = "None"
 CONST_NODE_PROVISIONING_DEFAULT_POOLS_AUTO = "Auto"
 
+# consts for workloadruntime
+CONST_WORKLOAD_RUNTIME_KATA_VM_ISOLATION = "KataVmIsolation"
+
 
 # consts for decorator pattern
 class DecoratorMode(Enum):

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

@@ -623,6 +623,13 @@
         Auto: A standard set of Karpenter NodePools are provisioned.
         None: No Karpenter NodePools are provisioned.
         WARNING: Changing this from Auto to None on an existing cluster will cause the default Karpenter NodePools to be deleted, which will in turn drain and delete the nodes associated with those pools. It is strongly recommended to not do this unless there are idle nodes ready to take the pods evicted by that action.
+  - name: --workload-runtime
+    type: string
+    short-summary: Set the workload runtime.
+    long-summary: |
+        Azure provides a different workload-runtime to enable Kata supported workloads in your nodepools. The following values can be specified:
+          - "KataVmIsolation" for Kata.
+
 examples:
   - name: Create a Kubernetes cluster with an existing SSH public key.
     text: az aks create -g MyResourceGroup -n MyManagedCluster --ssh-key-value /path/to/publickey
@@ -706,6 +713,8 @@
     text: az aks create -g MyResourceGroup -n MyManagedCluster --node-provisioning-mode Auto
   - name: Create a kubernetes cluster with auto node provisioning and no default pools.
     text: az aks create -g MyResourceGroup -n MyManagedCluster --node-provisioning-mode Auto --node-provisioning-default-pools None
+  - name: Create a Kubernetes cluster with KataVmIsolation enabled.
+    text: az aks create -g MyResourceGroup -n MyManagedCluster --os-sku AzureLinux --vm-size Standard_D4s_v3 --workload-runtime KataVmIsolation --node-count 1
 """
 
 helps["aks update"] = """
@@ -1953,6 +1962,13 @@
   - name: --gateway-prefix-size
     type: int
     short-summary: The size of Public IPPrefix attached to the Gateway-mode node pool. The node pool must be in Gateway mode.
+  - name: --workload-runtime
+    type: string
+    short-summary: Set the workload runtime.
+    long-summary: |
+        Azure provides a different workload-runtime to enable Kata supported workloads in your nodepools. The following values can be specified:
+          - "KataVmIsolation" for Kata.
+
 examples:
   - name: Create a nodepool in an existing AKS cluster with ephemeral os enabled.
     text: az aks nodepool add -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster --node-osdisk-type Ephemeral --node-osdisk-size 48
@@ -1974,6 +1990,8 @@
     text: az aks nodepool add -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster  --os-sku Ubuntu --pod-subnet-id /subscriptions/SubID/resourceGroups/AnotherResourceGroup/providers/Microsoft.Network/virtualNetworks/MyVnet/subnets/MySubnet --pod-ip-allocation-mode StaticBlock
   - name: create a nodepool of type VirtualMachines
     text: az aks nodepool add -g MyResourceGroup -n MyNodePool --cluster-name MyMC --vm-set-type VirtualMachines --vm-sizes "VMSize1,VMSize2" --node-count 3
+  - name: Create a nodepool with KataVmIsolation enabled.
+    text: az aks nodepool add -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster --os-sku AzureLinux --vm-size Standard_D4s_v3 --workload-runtime KataVmIsolation --node-count 1
 """
 
 helps["aks nodepool delete"] = """

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

@@ -73,7 +73,8 @@
     CONST_NODE_PROVISIONING_MODE_MANUAL,
     CONST_NODE_PROVISIONING_MODE_AUTO,
     CONST_NODE_PROVISIONING_DEFAULT_POOLS_NONE,
-    CONST_NODE_PROVISIONING_DEFAULT_POOLS_AUTO)
+    CONST_NODE_PROVISIONING_DEFAULT_POOLS_AUTO,
+    CONST_WORKLOAD_RUNTIME_KATA_VM_ISOLATION)
 from azure.cli.command_modules.acs.azurecontainerstorage._consts import (
     CONST_ACSTOR_ALL,
     CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY,
@@ -356,6 +357,10 @@
     CONST_APP_ROUTING_NONE_NGINX
 ]
 
+workload_runtime_types = [
+    CONST_WORKLOAD_RUNTIME_KATA_VM_ISOLATION,
+]
+
 
 def load_arguments(self, _):
     acr_arg_type = CLIArgumentType(metavar='ACR_NAME_OR_RESOURCE_ID')
@@ -536,6 +541,7 @@ def load_arguments(self, _):
         c.argument('gpu_instance_profile', arg_type=get_enum_type(gpu_instance_profiles))
         c.argument('nodepool_allowed_host_ports', nargs='+', validator=validate_allowed_host_ports, help="allowed host ports for agentpool")
         c.argument('nodepool_asg_ids', nargs='+', validator=validate_application_security_groups, help="application security groups for agentpool")
+        c.argument('workload_runtime', arg_type=get_enum_type(workload_runtime_types), help="The workload runtime to use on the node pool.")
         c.argument("message_of_the_day")
 
         # azure monitor profile
@@ -1046,6 +1052,7 @@ def load_arguments(self, _):
         c.argument("if_none_match")
         c.argument('gpu_driver', arg_type=get_enum_type(gpu_driver_install_modes))
         c.argument("gateway_prefix_size", type=int, validator=validate_gateway_prefix_size)
+        c.argument('workload_runtime', arg_type=get_enum_type(workload_runtime_types), help="The workload runtime to use on the nodepool.")
 
     with self.argument_context('aks nodepool update', resource_type=ResourceType.MGMT_CONTAINERSERVICE, operation_group='agent_pools') as c:
         c.argument('enable_cluster_autoscaler', options_list=[

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

@@ -1685,6 +1685,21 @@ def get_gateway_prefix_size(self) -> Union[int, None]:
         """
         return self.raw_param.get('gateway_prefix_size')
 
+    def get_workload_runtime(self) -> Union[str, None]:
+        """Obtain the value of workload_runtime, default value is None.
+
+        :return: string or None
+        """
+        # read the original value passed by the command
+        workload_runtime = self.raw_param.get("workload_runtime", None)
+        # try to read the property value corresponding to the parameter from the `mc` object
+        if self.agentpool and self.agentpool.workload_runtime is not None:
+            workload_runtime = self.agentpool.workload_runtime
+
+        # this parameter does not need dynamic completion
+        # this parameter does not need validation
+        return workload_runtime
+
 
 class AKSAgentPoolAddDecorator:
     def __init__(
@@ -2134,6 +2149,19 @@ def set_up_virtual_machines_profile(self, agentpool: AgentPool) -> AgentPool:
 
         return agentpool
 
+    def set_up_workload_runtime(self, agentpool: AgentPool) -> AgentPool:
+        """Set up workload runtime for the AgentPool object.
+
+        :return: the AgentPool object
+        """
+        self._ensure_agentpool(agentpool)
+
+        workload_runtime = self.context.get_workload_runtime()
+        if workload_runtime is not None:
+            agentpool.workload_runtime = workload_runtime
+
+        return agentpool
+
     def construct_agentpool_profile_default(self, bypass_restore_defaults: bool = False) -> AgentPool:
         """The overall controller used to construct the AgentPool profile by default.
 
@@ -2186,6 +2214,8 @@ def construct_agentpool_profile_default(self, bypass_restore_defaults: bool = Fa
         agentpool = self.set_up_agentpool_gateway_profile(agentpool)
         # set up virtual machines profile
         agentpool = self.set_up_virtual_machines_profile(agentpool)
+        # set up workload_runtime
+        agentpool = self.set_up_workload_runtime(agentpool)
         # restore defaults
         if not bypass_restore_defaults:
             agentpool = self._restore_defaults_in_agentpool(agentpool)

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

@@ -945,6 +945,7 @@ def aks_create(
     crg_id=None,
     gpu_instance_profile=None,
     message_of_the_day=None,
+    workload_runtime=None,
     # azure service mesh
     enable_azure_service_mesh=None,
     revision=None,
@@ -2855,6 +2856,7 @@ def aks_agentpool_add(
     asg_ids=None,
     node_public_ip_tags=None,
     disable_windows_outbound_nat=False,
+    workload_runtime=None,
     # trusted launch
     enable_vtpm=False,
     enable_secure_boot=False,