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.
 
 22.0.0b10
 +++++++++

src/aks-preview/README.rst

@@ -14,6 +14,29 @@ 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.
+Flatcar, Ubuntu 26.04, and Azure Linux 3 on ARM64 node pools are not supported
+by this extension version. The cluster must have at least one compatible node pool.
+Enabling an existing installation leaves its configuration unchanged.
+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

@@ -1292,6 +1292,13 @@ def load_arguments(self, _):
                    is_preview=True,
                    action="store_true"
                    )
+        c.argument(
+            "enable_azure_managed_lustre",
+            options_list=["--enable-azure-managed-lustre", "--enable-amlfs"],
+            action="store_true",
+            is_preview=True,
+            help="Install the Azure Managed Lustre extension. Requires the k8s-extension CLI extension.",
+        )
         # azure container storage
         c.argument(
             "enable_azure_container_storage",
@@ -2040,6 +2047,20 @@ 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",
+            options_list=["--enable-azure-managed-lustre", "--enable-amlfs"],
+            action="store_true",
+            is_preview=True,
+            help="Install the Azure Managed Lustre extension. Requires the k8s-extension CLI extension.",
+        )
+        c.argument(
+            "disable_azure_managed_lustre",
+            options_list=["--disable-azure-managed-lustre", "--disable-amlfs"],
+            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,11 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+CONST_AMLFS_EXT_INSTALLATION_NAME = "azurelustre"
+CONST_AMLFS_K8S_EXTENSION_NAME = "microsoft.azuremanagedlustre"
+CONST_AMLFS_VERSION = "0.6.0"
+CONST_AMLFS_RELEASE_TRAIN = "stable"
+CONST_AMLFS_UNSUPPORTED_OS_SKUS = frozenset(("flatcar", "ubuntu2604"))
+CONST_AMLFS_ARM64_UNSUPPORTED_OS_SKUS = frozenset(("azurelinux3", "ubuntu2204"))

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

@@ -0,0 +1,42 @@
+# --------------------------------------------------------------------------------------------
+# 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_AMLFS_EXT_INSTALLATION_NAME,
+    CONST_AMLFS_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 get_installed_extension(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_AMLFS_EXT_INSTALLATION_NAME, "managedClusters"
+        )
+    except ResourceNotFoundError:
+        return None
+
+    if extension.extension_type.lower() != CONST_AMLFS_K8S_EXTENSION_NAME:
+        raise InvalidArgumentValueError(
+            f"The extension '{CONST_AMLFS_EXT_INSTALLATION_NAME}' already exists with type "
+            f"'{extension.extension_type}', not '{CONST_AMLFS_K8S_EXTENSION_NAME}'."
+        )
+    return extension
+
+
+def check_if_extension_is_installed(cmd, resource_group, cluster_name):
+    return get_installed_extension(cmd, resource_group, cluster_name) is not None

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

@@ -0,0 +1,91 @@
+# --------------------------------------------------------------------------------------------
+# 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
+
+from azext_aks_preview.azuremanagedlustre._consts import (
+    CONST_AMLFS_ARM64_UNSUPPORTED_OS_SKUS,
+    CONST_AMLFS_UNSUPPORTED_OS_SKUS,
+)
+
+
+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."
+        )
+
+
+def _get_vm_sku_architectures(cli_ctx, location):
+    from azure.cli.command_modules.vm.aaz.latest.vm import ListSkus as VMListSkus
+
+    def is_in_location(sku):
+        return not location or any(location.lower() == value.lower() for value in sku.get("locations", []))
+
+    architectures = {}
+    for sku in VMListSkus(cli_ctx=cli_ctx)(command_args={}):
+        if sku.get("resourceType", "").lower() != "virtualmachines" or not is_in_location(sku):
+            continue
+        for capability in sku.get("capabilities", []):
+            if capability.get("name", "").lower() == "cpuarchitecturetype":
+                architectures[sku.get("name", "").lower()] = capability.get("value", "").lower()
+                break
+    return architectures
+
+
+def _resolves_to_azure_linux_3(os_sku, kubernetes_version):
+    if os_sku != "azurelinux" or not kubernetes_version:
+        return False
+    version_parts = str(kubernetes_version).lstrip("v").split(".")
+    try:
+        return (int(version_parts[0]), int(version_parts[1].split("-")[0])) >= (1, 32)
+    except (IndexError, ValueError):
+        return False
+
+
+def validate_azure_managed_lustre_node_compatibility(cmd, cluster):
+    pools = cluster.agent_pool_profiles or []
+    architecture_sensitive_pools = []
+
+    for pool in pools:
+        os_sku = (pool.os_sku or "").lower()
+        if os_sku in CONST_AMLFS_UNSUPPORTED_OS_SKUS:
+            continue
+        kubernetes_version = (
+            getattr(pool, "orchestrator_version", None) or getattr(cluster, "kubernetes_version", None)
+        )
+        if (
+            os_sku in CONST_AMLFS_ARM64_UNSUPPORTED_OS_SKUS or
+            _resolves_to_azure_linux_3(os_sku, kubernetes_version)
+        ):
+            architecture_sensitive_pools.append(pool)
+            continue
+        return
+
+    if architecture_sensitive_pools:
+        architectures = _get_vm_sku_architectures(cmd.cli_ctx, cluster.location)
+        unknown_vm_sizes = []
+        for pool in architecture_sensitive_pools:
+            vm_size = (pool.vm_size or "").lower()
+            architecture = architectures.get(vm_size)
+            if architecture and architecture != "arm64":
+                return
+            if not architecture:
+                unknown_vm_sizes.append(pool.vm_size or "<unknown>")
+        if unknown_vm_sizes:
+            raise InvalidArgumentValueError(
+                "Cannot enable Azure Managed Lustre because the architecture couldn't be determined for "
+                f"node pool VM size(s): {', '.join(unknown_vm_sizes)}."
+            )
+
+    raise InvalidArgumentValueError(
+        "Cannot enable Azure Managed Lustre because none of the cluster's node pools support "
+        "Azure Managed Lustre extension version 0.6.0. Flatcar and Ubuntu 26.04 aren't supported; "
+        "Ubuntu 22.04 and Azure Linux 3 on ARM64 aren't supported."
+    )

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_AMLFS_EXT_INSTALLATION_NAME,
+    CONST_AMLFS_K8S_EXTENSION_NAME,
+    CONST_AMLFS_RELEASE_TRAIN,
+    CONST_AMLFS_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_AMLFS_EXT_INSTALLATION_NAME,
+        "managedClusters",
+        CONST_AMLFS_K8S_EXTENSION_NAME,
+        version=CONST_AMLFS_VERSION,
+        release_train=CONST_AMLFS_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_AMLFS_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,