Add AMLFS update preview command

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
 +++++++
+* `az aks create/update`: Add `--enable-azure-managed-lustre` to install the Azure Managed Lustre cluster extension at version 0.6.0 on the stable release train with automatic upgrades disabled. Add `az aks update --disable-azure-managed-lustre` to uninstall it.
 * `az aks nodepool update`: Preserve the existing GPU management mode when `--enable-managed-gpu` is omitted, including when enabling, updating, or disabling the cluster autoscaler.
 * `az aks alert-config add`: Reject an empty `--name` before looking up existing configurations instead of reporting that it already exists.
 * `az aks nodepool scale`: add `--use-patch-api` to optionally scale a VMSS node pool via the new dedicated PATCH agent pool API (scales to the target count without triggering full reconciliation). The default behavior continues to use the PUT agent pool API.

src/aks-preview/README.rst

@@ -14,6 +14,26 @@ Install this extension using the following CLI command `az extension add --name
 
 Remove this extension using the following CLI command `az extension remove --name aks-preview`.
 
+Azure Managed Lustre
+--------------------
+
+Install the ``k8s-extension`` CLI extension before enabling Azure Managed Lustre::
+
+    az extension add --name k8s-extension
+    az aks create --resource-group <resource-group> --name <cluster-name> --generate-ssh-keys --enable-azure-managed-lustre
+
+For an existing AKS cluster, enable or disable the extension with::
+
+    az aks update --resource-group <resource-group> --name <cluster-name> --enable-azure-managed-lustre
+    az aks update --resource-group <resource-group> --name <cluster-name> --disable-azure-managed-lustre
+
+Enabling installs ``microsoft.azuremanagedlustre`` as ``azurelustre`` with cluster scope,
+version ``0.6.0``, release train ``stable``, and automatic upgrades disabled.
+As with Azure Container Storage, these operations wait for cluster provisioning before
+installing or uninstalling the extension, even when ``--no-wait`` is specified.
+The enable and disable flags cannot be used together. Disabling uninstalls the cluster
+extension; it does not delete an Azure Managed Lustre file system.
+
 Dependency between aks-preview and azure-cli/acs (azure-cli-core)
 =================================================================
 

src/aks-preview/azext_aks_preview/_params.py

@@ -1289,6 +1289,13 @@ def load_arguments(self, _):
                    is_preview=True,
                    action="store_true"
                    )
+        c.argument(
+            "enable_azure_managed_lustre",
+            action="store_true",
+            is_preview=True,
+            help="Install the Azure Managed Lustre extension version 0.6.0 on the stable release train "
+                 "with automatic upgrades disabled. Requires the k8s-extension CLI extension.",
+        )
         # azure container storage
         c.argument(
             "enable_azure_container_storage",
@@ -2037,6 +2044,19 @@ def load_arguments(self, _):
         c.argument("disable_cost_analysis", action="store_true")
         c.argument('enable_ai_toolchain_operator', is_preview=True, action='store_true')
         c.argument('disable_ai_toolchain_operator', is_preview=True, action='store_true')
+        c.argument(
+            "enable_azure_managed_lustre",
+            action="store_true",
+            is_preview=True,
+            help="Install the Azure Managed Lustre extension version 0.6.0 on the stable release train "
+                 "with automatic upgrades disabled. Requires the k8s-extension CLI extension.",
+        )
+        c.argument(
+            "disable_azure_managed_lustre",
+            action="store_true",
+            is_preview=True,
+            help="Uninstall the Azure Managed Lustre extension from the cluster.",
+        )
         # azure container storage
         c.argument(
             "enable_azure_container_storage",

src/aks-preview/azext_aks_preview/azuremanagedlustre/__init__.py

@@ -0,0 +1,4 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------

src/aks-preview/azext_aks_preview/azuremanagedlustre/_consts.py

@@ -0,0 +1,9 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+CONST_AML_EXT_INSTALLATION_NAME = "azurelustre"
+CONST_AML_K8S_EXTENSION_NAME = "microsoft.azuremanagedlustre"
+CONST_AML_VERSION = "0.6.0"
+CONST_AML_RELEASE_TRAIN = "stable"

src/aks-preview/azext_aks_preview/azuremanagedlustre/_helpers.py

@@ -0,0 +1,38 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from azure.cli.core.azclierror import InvalidArgumentValueError, ResourceNotFoundError
+from azext_aks_preview._consts import (
+    CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME,
+    CONST_K8S_EXTENSION_CUSTOM_MOD_NAME,
+)
+from azext_aks_preview._helpers import get_k8s_extension_module
+from azext_aks_preview.azuremanagedlustre._consts import (
+    CONST_AML_EXT_INSTALLATION_NAME,
+    CONST_AML_K8S_EXTENSION_NAME,
+)
+
+
+def get_azure_managed_lustre_extension_client(cmd):
+    client_factory = get_k8s_extension_module(CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME)
+    custom_module = get_k8s_extension_module(CONST_K8S_EXTENSION_CUSTOM_MOD_NAME)
+    return client_factory.cf_k8s_extension_operation(cmd.cli_ctx), custom_module
+
+
+def check_if_extension_is_installed(cmd, resource_group, cluster_name):
+    client, custom_module = get_azure_managed_lustre_extension_client(cmd)
+    try:
+        extension = custom_module.show_k8s_extension(
+            client, resource_group, cluster_name, CONST_AML_EXT_INSTALLATION_NAME, "managedClusters"
+        )
+    except ResourceNotFoundError:
+        return False
+
+    if extension.extension_type.lower() != CONST_AML_K8S_EXTENSION_NAME:
+        raise InvalidArgumentValueError(
+            f"The extension '{CONST_AML_EXT_INSTALLATION_NAME}' already exists with type "
+            f"'{extension.extension_type}', not '{CONST_AML_K8S_EXTENSION_NAME}'."
+        )
+    return True

src/aks-preview/azext_aks_preview/azuremanagedlustre/_validators.py

@@ -0,0 +1,17 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from azure.cli.core.azclierror import InvalidArgumentValueError, MutuallyExclusiveArgumentError
+
+
+def validate_azure_managed_lustre_params(enable, disable, is_extension_installed=None):
+    if enable and disable:
+        raise MutuallyExclusiveArgumentError(
+            "Cannot set --enable-azure-managed-lustre and --disable-azure-managed-lustre together."
+        )
+    if disable and is_extension_installed is False:
+        raise InvalidArgumentValueError(
+            "Cannot set --disable-azure-managed-lustre. Azure Managed Lustre is not enabled in the cluster."
+        )

src/aks-preview/azext_aks_preview/azuremanagedlustre/aml_ops.py

@@ -0,0 +1,57 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from azure.cli.core.azclierror import UnknownError
+from azure.cli.core.commands import LongRunningOperation
+from knack.log import get_logger
+
+from azext_aks_preview.azuremanagedlustre._consts import (
+    CONST_AML_EXT_INSTALLATION_NAME,
+    CONST_AML_K8S_EXTENSION_NAME,
+    CONST_AML_RELEASE_TRAIN,
+    CONST_AML_VERSION,
+)
+from azext_aks_preview.azuremanagedlustre._helpers import get_azure_managed_lustre_extension_client
+
+logger = get_logger(__name__)
+
+
+def perform_enable_azure_managed_lustre(cmd, resource_group, cluster_name):
+    client, custom_module = get_azure_managed_lustre_extension_client(cmd)
+    result = custom_module.create_k8s_extension(
+        cmd,
+        client,
+        resource_group,
+        cluster_name,
+        CONST_AML_EXT_INSTALLATION_NAME,
+        "managedClusters",
+        CONST_AML_K8S_EXTENSION_NAME,
+        version=CONST_AML_VERSION,
+        release_train=CONST_AML_RELEASE_TRAIN,
+        scope="cluster",
+        auto_upgrade_minor_version=False,
+    )
+    extension = LongRunningOperation(cmd.cli_ctx)(result)
+    if extension is None or extension.provisioning_state != "Succeeded":
+        raise UnknownError("Azure Managed Lustre extension installation did not succeed.")
+    logger.warning("Azure Managed Lustre successfully installed.")
+
+
+def perform_disable_azure_managed_lustre(cmd, resource_group, cluster_name):
+    client, custom_module = get_azure_managed_lustre_extension_client(cmd)
+    result = custom_module.delete_k8s_extension(
+        cmd,
+        client,
+        resource_group,
+        cluster_name,
+        CONST_AML_EXT_INSTALLATION_NAME,
+        "managedClusters",
+        yes=True,
+    )
+    # The delegated delete returns None when it cannot retrieve the extension.
+    if result is None:
+        raise UnknownError("Azure Managed Lustre extension could not be retrieved for deletion.")
+    LongRunningOperation(cmd.cli_ctx)(result)
+    logger.warning("Azure Managed Lustre has been disabled.")

src/aks-preview/azext_aks_preview/custom.py

@@ -1440,6 +1440,7 @@ def aks_create(
     enable_cost_analysis=False,
     # AI toolchain operator
     enable_ai_toolchain_operator=False,
+    enable_azure_managed_lustre=False,
     # azure container storage
     enable_azure_container_storage=None,
     container_storage_version=None,
@@ -1714,6 +1715,8 @@ def aks_update(
     # AI toolchain operator
     enable_ai_toolchain_operator=False,
     disable_ai_toolchain_operator=False,
+    enable_azure_managed_lustre=False,
+    disable_azure_managed_lustre=False,
     # azure container storage
     enable_azure_container_storage=None,
     disable_azure_container_storage=None,

src/aks-preview/azext_aks_preview/managed_cluster_decorator.py

@@ -96,6 +96,15 @@
     perform_enable_azure_container_storage_v1,
     perform_azure_container_storage_update,
 )
+from azext_aks_preview.azuremanagedlustre._helpers import (
+    check_if_extension_is_installed as check_if_azure_managed_lustre_is_installed,
+    get_azure_managed_lustre_extension_client,
+)
+from azext_aks_preview.azuremanagedlustre._validators import validate_azure_managed_lustre_params
+from azext_aks_preview.azuremanagedlustre.aml_ops import (
+    perform_disable_azure_managed_lustre,
+    perform_enable_azure_managed_lustre,
+)
 from azext_aks_preview.azuremonitormetrics.azuremonitorprofile import (
     ensure_azure_monitor_profile_prerequisites,
 )
@@ -291,6 +300,8 @@ def external_functions(self) -> SimpleNamespace:
                 "perform_disable_azure_container_storage_v1"
             ] = perform_disable_azure_container_storage_v1
             external_functions["perform_azure_container_storage_update"] = perform_azure_container_storage_update
+            external_functions["perform_enable_azure_managed_lustre"] = perform_enable_azure_managed_lustre
+            external_functions["perform_disable_azure_managed_lustre"] = perform_disable_azure_managed_lustre
             external_functions["sanitize_loganalytics_ws_resource_id"] = sanitize_loganalytics_ws_resource_id
             # Override base module function with preview version that uses REST API to avoid
             # "Request Header Fields Too Large" errors
@@ -5358,6 +5369,13 @@ def set_up_azure_monitor_profile(self, mc: ManagedCluster) -> ManagedCluster:
 
         return mc
 
+    def set_up_azure_managed_lustre(self, mc: ManagedCluster) -> ManagedCluster:
+        self._ensure_mc(mc)
+        if self.context.raw_param.get("enable_azure_managed_lustre"):
+            get_azure_managed_lustre_extension_client(self.cmd)
+            self.context.set_intermediate("enable_azure_managed_lustre", True, overwrite_exists=True)
+        return mc
+
     def set_up_azure_container_storage(self, mc: ManagedCluster) -> ManagedCluster:  # pylint: disable=too-many-locals
         """Set up azure container storage for the Managed Cluster object
         :return: ManagedCluster
@@ -5952,6 +5970,7 @@ def construct_mc_profile_preview(self, bypass_restore_defaults: bool = False) ->
         mc = self.set_up_ai_toolchain_operator(mc)
         # set up for azure container storage
         mc = self.set_up_azure_container_storage(mc)
+        mc = self.set_up_azure_managed_lustre(mc)
         # set up node provisioning profile
         mc = self.set_up_node_provisioning_profile(mc)
         # set up node disruption policy
@@ -6044,6 +6063,7 @@ def check_is_postprocessing_required(self, mc: ManagedCluster) -> bool:
             (enable_managed_identity and attach_acr) or
             need_grant_vnet_permission_to_cluster_identity or
             enable_azure_container_storage or
+            self.context.get_intermediate("enable_azure_managed_lustre", default_value=False) or
             enable_backup
         ):
             return True
@@ -6303,6 +6323,11 @@ def postprocessing_after_mc_created(self, cluster: ManagedCluster) -> None:
                     is_called_from_extension=True,
                 )
 
+        if self.context.get_intermediate("enable_azure_managed_lustre", default_value=False):
+            self.context.external_functions.perform_enable_azure_managed_lustre(
+                self.cmd, self.context.get_resource_group_name(), self.context.get_name()
+            )
+
         # Add role assignments for automatic sku
         if cluster.sku is not None and cluster.sku.name == "Automatic":
             try:
@@ -6733,6 +6758,20 @@ def update_monitoring_profile_flow_logs(self, mc: ManagedCluster) -> ManagedClus
 
         return mc
 
+    def update_azure_managed_lustre(self, mc: ManagedCluster) -> ManagedCluster:
+        self._ensure_mc(mc)
+        enable = self.context.raw_param.get("enable_azure_managed_lustre", False)
+        disable = self.context.raw_param.get("disable_azure_managed_lustre", False)
+        validate_azure_managed_lustre_params(enable, disable)
+        if enable or disable:
+            installed = check_if_azure_managed_lustre_is_installed(
+                self.cmd, self.context.get_resource_group_name(), self.context.get_name()
+            )
+            validate_azure_managed_lustre_params(enable, disable, installed)
+            self.context.set_intermediate("enable_azure_managed_lustre", enable, overwrite_exists=True)
+            self.context.set_intermediate("disable_azure_managed_lustre", disable, overwrite_exists=True)
+        return mc
+
     # pylint: disable=too-many-statements,too-many-locals,too-many-branches
     def update_azure_container_storage(self, mc: ManagedCluster) -> ManagedCluster:
         """Update azure container storage for the Managed Cluster object
@@ -9106,6 +9145,7 @@ def update_mc_profile_preview(self) -> ManagedCluster:
         mc = self.update_ai_toolchain_operator(mc)
         # update azure container storage
         mc = self.update_azure_container_storage(mc)
+        mc = self.update_azure_managed_lustre(mc)
         # update node provisioning profile
         mc = self.update_node_provisioning_profile(mc)
         # update bootstrap profile
@@ -9157,6 +9197,8 @@ def check_is_postprocessing_required(self, mc: ManagedCluster) -> bool:
             # Note: monitoring_addon_disable_postprocessing_required is no longer used - cleanup is done upfront
             # pylint: disable=too-many-boolean-expressions
             if (enable_azure_container_storage or disable_azure_container_storage) or \
+               self.context.get_intermediate("enable_azure_managed_lustre", default_value=False) or \
+               self.context.get_intermediate("disable_azure_managed_lustre", default_value=False) or \
                (keyvault_id and enable_azure_keyvault_secrets_provider_addon) or \
                (monitoring_addon_postprocessing_required) or \
                enable_backup:
@@ -9172,6 +9214,15 @@ def postprocessing_after_mc_created(self, cluster: ManagedCluster) -> None:
         """
         super().postprocessing_after_mc_created(cluster)
 
+        if self.context.get_intermediate("enable_azure_managed_lustre", default_value=False):
+            self.context.external_functions.perform_enable_azure_managed_lustre(
+                self.cmd, self.context.get_resource_group_name(), self.context.get_name()
+            )
+        if self.context.get_intermediate("disable_azure_managed_lustre", default_value=False):
+            self.context.external_functions.perform_disable_azure_managed_lustre(
+                self.cmd, self.context.get_resource_group_name(), self.context.get_name()
+            )
+
         # Handle monitoring addon postprocessing (enable case)
         monitoring_addon_postprocessing_required = self.context.get_intermediate(
             "monitoring_addon_postprocessing_required", default_value=False