Adding enable ultra ssd, eviction policy, and spot max price support for machine (#9696)

Author: erinborders

src/aks-preview/HISTORY.rst

@@ -11,6 +11,9 @@ To release a new version, please select a new version number (usually plus 1 to
 
 Pending
 +++++++
+* `az aks machine add`: Add `--spot-max-price` flag support to set the max price (in US Dollars) you are willing to pay for spot instances on a machine.
+* `az aks machine add`: Add `--eviction-policy` flag support to set the eviction policy for a machine.
+* `az aks machine add`: Add `--enable-ultra-ssd` flag support to enable ultra ssd on a machine.
 
 20.0.0b1
 +++++++

src/aks-preview/azext_aks_preview/_help.py

@@ -2731,6 +2731,15 @@
        - name: --node-public-ip-tags
          type: string
          short-summary: The ipTags of the machine public IPs.
+       - name: --spot-max-price
+         type: number
+         short-summary: The max price (in US Dollars) you are willing to pay for spot instances.
+       - name: --enable-ultra-ssd
+         type: bool
+         short-summary: Whether to enable UltraSSD.
+       - name: --eviction-policy
+         type: string
+         short-summary: The eviction policy for machine. This cannot be specified unless the priority is 'Spot'. If not specified, the default is 'Delete'.
 """
 
 helps['aks machine update'] = """

src/aks-preview/azext_aks_preview/_params.py

@@ -2438,6 +2438,21 @@ def load_arguments(self, _):
             validator=validate_k8s_version,
             help="Version of Kubernetes to use for the machine.",
         )
+        c.argument(
+            "spot_max_price",
+            type=float,
+            validator=validate_spot_max_price,
+            help="The max price (in US Dollars) you are willing to pay for spot instances."
+        )
+        c.argument(
+            "enable_ultra_ssd",
+            action="store_true"
+        )
+        c.argument(
+            "eviction_policy",
+            arg_type=get_enum_type(node_eviction_policies),
+            validator=validate_eviction_policy,
+        )
 
     with self.argument_context("aks machine update") as c:
         c.argument(

src/aks-preview/azext_aks_preview/custom.py

@@ -15,6 +15,7 @@
 import time
 import webbrowser
 import subprocess
+from math import isnan
 
 from azext_aks_preview._client_factory import (
     CUSTOM_MGMT_AKS_PREVIEW,
@@ -2761,6 +2762,9 @@ def aks_machine_add(
     vm_size=None,
     kubernetes_version=None,
     no_wait=False,
+    spot_max_price=float("nan"),
+    enable_ultra_ssd=False,
+    eviction_policy=None,
 ):
     existedMachine = None
     try:
@@ -2773,6 +2777,9 @@ def aks_machine_add(
             f"Machine '{machine_name}' already exists. Please use 'az aks machine update' to update it."
         )
 
+    if isnan(spot_max_price):
+        spot_max_price = -1
+
     # DO NOT MOVE: get all the original parameters and save them as a dictionary
     raw_parameters = locals()
     return add_machine(cmd, client, raw_parameters, no_wait)

src/aks-preview/azext_aks_preview/machine.py

@@ -84,14 +84,18 @@ def constructMachine(cmd, raw_parameters, machine_name):
     )
     tags = raw_parameters.get("tags")
     priority = raw_parameters.get("priority")
+    eviction_policy = raw_parameters.get("eviction_policy")
     machineProperties = MachineProperties(
         tags=tags,
         priority=priority,
+        eviction_policy=eviction_policy,
         network=set_machine_network(cmd, raw_parameters),
         hardware=set_machine_hardware_profile(cmd, raw_parameters),
         kubernetes=set_machine_kubernetes_profile(cmd, raw_parameters),
-        operating_system=set_machine_os_profile(cmd, raw_parameters)
+        operating_system=set_machine_os_profile(cmd, raw_parameters),
+        billing=set_machine_billing_profile(cmd, raw_parameters)
     )
+
     Machine = cmd.get_models(
         "Machine",
         resource_type=CUSTOM_MGMT_AKS_PREVIEW,
@@ -106,6 +110,7 @@ def constructMachine(cmd, raw_parameters, machine_name):
 
 
 def set_machine_hardware_profile(cmd, raw_parameters):
+    enable_ultra_ssd = raw_parameters.get("enable_ultra_ssd")
     vm_size = raw_parameters.get("vm_size")
     if vm_size is None:
         raise RequiredArgumentMissingError(
@@ -117,7 +122,8 @@ def set_machine_hardware_profile(cmd, raw_parameters):
         operation_group="machines"
     )
     machine_hardware_profile = MachineHardwareProfile(
-        vm_size=vm_size
+        vm_size=vm_size,
+        ultra_ssd_enabled=enable_ultra_ssd
     )
     return machine_hardware_profile
 
@@ -196,3 +202,16 @@ def set_machine_os_profile(cmd, raw_parameters):
         enable_fips=enable_fips
     )
     return machineOSProfile
+
+
+def set_machine_billing_profile(cmd, raw_parameters):
+    spot_max_price = raw_parameters.get("spot_max_price")
+    MachineBillingProfile = cmd.get_models(
+        "MachineBillingProfile",
+        resource_type=CUSTOM_MGMT_AKS_PREVIEW,
+        operation_group="machines"
+    )
+    machineBillingProfile = MachineBillingProfile(
+        spot_max_price=spot_max_price
+    )
+    return machineBillingProfile