{AKS} az aks nodepool add/update/upgrade: add new parameter --node-soak-duration

src/aks-preview/HISTORY.rst

@@ -13,6 +13,11 @@ Pending
 +++++++
 * Vendor new SDK and bump API version to 2023-09-02-preview.
 
+0.5.169
++++++++
+* Add `--node-soak-duration` to the `az aks nodepool add/update/upgrade` commands.
+* Add `--drain-timeout` to the `az aks nodepool add/update/upgrade` commands (already in [azure-cli](https://github.com/Azure/azure-cli/pull/27475)).
+
 0.5.168
 +++++++
 * Add `--enable-image-integrity` to the `az aks update` command.

src/aks-preview/azext_aks_preview/_help.py

@@ -1610,6 +1610,12 @@
         - name: --max-surge
           type: string
           short-summary: Extra nodes used to speed upgrade. When specified, it represents the number or percent used, eg. 5 or 33%
+        - name: --drain-timeout
+          type: int
+          short-summary: When nodes are drain how many minutes to wait for all pods to be evicted
+        - name: --node-soak-duration
+          type: int
+          short-summary: The amount of time (in minutes) to wait after draining a node and before reimaging it and moving on to next node.
         - name: --kubelet-config
           type: string
           short-summary: Kubelet configurations for agent nodes.
@@ -1697,6 +1703,12 @@
         - name: --max-surge
           type: string
           short-summary: Extra nodes used to speed upgrade. When specified, it represents the number or percent used, eg. 5 or 33%
+        - name: --drain-timeout
+          type: int
+          short-summary: When nodes are drain how many minutes to wait for all pods to be evicted
+        - name: --node-soak-duration
+          type: int
+          short-summary: The amount of time (in minutes) to wait after draining a node and before reimaging it and moving on to next node.
         - name: --aks-custom-headers
           type: string
           short-summary: Send custom headers. When specified, format should be Key1=Value1,Key2=Value2
@@ -1731,6 +1743,12 @@
         - name: --max-surge
           type: string
           short-summary: Extra nodes used to speed upgrade. When specified, it represents the number or percent used, eg. 5 or 33%
+        - name: --drain-timeout
+          type: int
+          short-summary: When nodes are drain how many minutes to wait for all pods to be evicted
+        - name: --node-soak-duration
+          type: int
+          short-summary: The amount of time (in minutes) to wait after draining a node and before reimaging it and moving on to next node.
         - name: --mode
           type: string
           short-summary: The mode for a node pool which defines a node pool's primary function. If set as "System", AKS prefers system pods scheduling to node pools with mode `System`. Learn more at https://aka.ms/aks/nodepool/mode.

src/aks-preview/azext_aks_preview/_params.py

@@ -685,6 +685,8 @@ def load_arguments(self, _):
         c.argument('node_osdisk_type', arg_type=get_enum_type(node_os_disk_types))
         c.argument('node_osdisk_size', type=int)
         c.argument('max_surge', validator=validate_max_surge)
+        c.argument('drain_timeout', type=int)
+        c.argument('node_soak_duration', type=int)
         c.argument('mode', arg_type=get_enum_type(node_mode_types))
         c.argument('scale_down_mode', arg_type=get_enum_type(scale_down_modes))
         c.argument('max_pods', type=int, options_list=['--max-pods', '-m'])
@@ -724,6 +726,8 @@ def load_arguments(self, _):
         c.argument('tags', tags_type)
         c.argument('node_taints', validator=validate_nodepool_taints)
         c.argument('max_surge', validator=validate_max_surge)
+        c.argument('drain_timeout', type=int)
+        c.argument('node_soak_duration', type=int)
         c.argument('mode', arg_type=get_enum_type(node_mode_types))
         c.argument('scale_down_mode', arg_type=get_enum_type(scale_down_modes))
         # extensions
@@ -734,6 +738,8 @@ def load_arguments(self, _):
 
     with self.argument_context('aks nodepool upgrade') as c:
         c.argument('max_surge', validator=validate_max_surge)
+        c.argument('drain_timeout', type=int)
+        c.argument('node_soak_duration', type=int)
         c.argument('snapshot_id', validator=validate_snapshot_id)
         c.argument('yes', options_list=['--yes', '-y'], help='Do not prompt for confirmation.', action='store_true')
         c.argument('aks_custom_headers')

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -327,6 +327,46 @@ def get_node_taints(self) -> Union[List[str], None]:
         # this parameter does not need validation
         return node_taints
 
+    def get_drain_timeout(self):
+        """Obtain the value of drain_timeout.
+
+        :return: int
+        """
+        # read the original value passed by the command
+        drain_timeout = self.raw_param.get("drain_timeout")
+        # 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.upgrade_settings and
+                self.agentpool.upgrade_settings.drain_timeout_in_minutes is not None
+            ):
+                drain_timeout = self.agentpool.upgrade_settings.drain_timeout_in_minutes
+
+        # this parameter does not need dynamic completion
+        # this parameter does not need validation
+        return drain_timeout
+
+    def get_node_soak_duration(self):
+        """Obtain the value of node_soak_duration.
+
+        :return: int
+        """
+        # read the original value passed by the command
+        node_soak_duration = self.raw_param.get("node_soak_duration")
+        # 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.upgrade_settings and
+                self.agentpool.upgrade_settings.node_soak_duration_in_minutes is not None
+            ):
+                node_soak_duration = self.agentpool.upgrade_settings.node_soak_duration_in_minutes
+
+        # this parameter does not need dynamic completion
+        # this parameter does not need validation
+        return node_soak_duration
+
 
 class AKSPreviewAgentPoolAddDecorator(AKSAgentPoolAddDecorator):
     def __init__(
@@ -472,6 +512,29 @@ def construct_agentpool_profile_preview(self) -> AgentPool:
         agentpool = self._restore_defaults_in_agentpool(agentpool)
         return agentpool
 
+    def set_up_upgrade_settings(self, agentpool: AgentPool) -> AgentPool:
+        """Set up upgrade settings for the AgentPool object.
+
+        :return: the AgentPool object
+        """
+        self._ensure_agentpool(agentpool)
+
+        upgrade_settings = self.models.AgentPoolUpgradeSettings()
+        max_surge = self.context.get_max_surge()
+        if max_surge:
+            upgrade_settings.max_surge = max_surge
+
+        drain_timeout = self.context.get_drain_timeout()
+        if drain_timeout:
+            upgrade_settings.drain_timeout_in_minutes = drain_timeout
+
+        node_soak_duration = self.context.get_node_soak_duration()
+        if node_soak_duration:
+            upgrade_settings.node_soak_duration_in_minutes = node_soak_duration
+
+        agentpool.upgrade_settings = upgrade_settings
+        return agentpool
+
 
 class AKSPreviewAgentPoolUpdateDecorator(AKSAgentPoolUpdateDecorator):
     def __init__(
@@ -551,3 +614,32 @@ def update_agentpool_profile_preview(self, agentpools: List[AgentPool] = None) -
         agentpool = self.update_network_profile(agentpool)
 
         return agentpool
+
+    def update_upgrade_settings(self, agentpool: AgentPool) -> AgentPool:
+        """Update upgrade settings for the Agentpool object.
+
+        :return: the Agentpool object
+        """
+        self._ensure_agentpool(agentpool)
+
+        upgrade_settings = agentpool.upgrade_settings
+        if upgrade_settings is None:
+            upgrade_settings = self.models.AgentPoolUpgradeSettings()
+
+        max_surge = self.context.get_max_surge()
+        if max_surge:
+            upgrade_settings.max_surge = max_surge
+            # why not always set this? so we don't wipe out a preview feaure in upgrade settigns like NodeSoakDuration?
+            agentpool.upgrade_settings = upgrade_settings
+
+        drain_timeout = self.context.get_drain_timeout()
+        if drain_timeout:
+            upgrade_settings.drain_timeout_in_minutes = drain_timeout
+            agentpool.upgrade_settings = upgrade_settings
+
+        node_soak_duration = self.context.get_node_soak_duration()
+        if node_soak_duration:
+            upgrade_settings.node_soak_duration_in_minutes = node_soak_duration
+            agentpool.upgrade_settings = upgrade_settings
+
+        return agentpool

src/aks-preview/azext_aks_preview/custom.py

@@ -1096,6 +1096,8 @@ def aks_agentpool_add(
     node_osdisk_type=None,
     node_osdisk_size=0,
     max_surge=None,
+    drain_timeout=None,
+    node_soak_duration=None,
     mode=CONST_NODEPOOL_MODE_USER,
     scale_down_mode=CONST_SCALE_DOWN_MODE_DELETE,
     max_pods=0,
@@ -1161,6 +1163,8 @@ def aks_agentpool_update(
     tags=None,
     node_taints=None,
     max_surge=None,
+    drain_timeout=None,
+    node_soak_duration=None,
     mode=None,
     scale_down_mode=None,
     no_wait=False,
@@ -1222,6 +1226,8 @@ def aks_agentpool_upgrade(cmd,
                           kubernetes_version='',
                           node_image_only=False,
                           max_surge=None,
+                          drain_timeout=None,
+                          node_soak_duration=None,
                           snapshot_id=None,
                           no_wait=False,
                           aks_custom_headers=None,
@@ -1239,11 +1245,11 @@ def aks_agentpool_upgrade(cmd,
         )
 
     # Note: we exclude this option because node image upgrade can't accept nodepool put fields like max surge
-    if max_surge and node_image_only:
+    if (max_surge or drain_timeout or node_soak_duration) and node_image_only:
         raise MutuallyExclusiveArgumentError(
-            'Conflicting flags. Unable to specify max-surge with node-image-only.'
-            'If you want to use max-surge with a node image upgrade, please first '
-            'update max-surge using "az aks nodepool update --max-surge".'
+            'Conflicting flags. Unable to specify max-surge/drain-timeout/node-soak-duration with node-image-only.'
+            'If you want to use max-surge/drain-timeout/node-soak-duration with a node image upgrade, please first '
+            'update max-surge/drain-timeout/node-soak-duration using "az aks nodepool update --max-surge/--drain-timeout/--node-soak-duration".'
         )
 
     if node_image_only:
@@ -1290,6 +1296,10 @@ def aks_agentpool_upgrade(cmd,
 
     if max_surge:
         instance.upgrade_settings.max_surge = max_surge
+    if drain_timeout:
+        instance.upgrade_settings.drain_timeout_in_minutes = drain_timeout
+    if node_soak_duration:
+        instance.upgrade_settings.node_soak_duration_in_minutes = node_soak_duration
 
     # custom headers
     aks_custom_headers = extract_comma_separated_string(