Fix Managed Lustre review findings

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: wolfgang-desalvador

src/aks-preview/README.rst

@@ -29,6 +29,9 @@ For an existing AKS cluster, enable or disable the extension with::
 
 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

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

@@ -7,3 +7,5 @@
 CONST_AML_K8S_EXTENSION_NAME = "microsoft.azuremanagedlustre"
 CONST_AML_VERSION = "0.6.0"
 CONST_AML_RELEASE_TRAIN = "stable"
+CONST_AML_UNSUPPORTED_OS_SKUS = frozenset(("flatcar", "ubuntu2604"))
+CONST_AML_ARM64_UNSUPPORTED_OS_SKUS = frozenset(("azurelinux3",))

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

@@ -5,6 +5,11 @@
 
 from azure.cli.core.azclierror import InvalidArgumentValueError, MutuallyExclusiveArgumentError
 
+from azext_aks_preview.azuremanagedlustre._consts import (
+    CONST_AML_ARM64_UNSUPPORTED_OS_SKUS,
+    CONST_AML_UNSUPPORTED_OS_SKUS,
+)
+
 
 def validate_azure_managed_lustre_params(enable, disable, is_extension_installed=None):
     if enable and disable:
@@ -15,3 +20,56 @@ def validate_azure_managed_lustre_params(enable, disable, is_extension_installed
         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 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_AML_UNSUPPORTED_OS_SKUS:
+            continue
+        if os_sku in CONST_AML_ARM64_UNSUPPORTED_OS_SKUS:
+            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, Ubuntu 26.04, and Azure Linux 3 "
+        "on ARM64 node pools aren't supported."
+    )

src/aks-preview/azext_aks_preview/managed_cluster_decorator.py

@@ -104,7 +104,10 @@
     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._validators import (
+    validate_azure_managed_lustre_node_compatibility,
+    validate_azure_managed_lustre_params,
+)
 from azext_aks_preview.azuremanagedlustre.aml_ops import (
     perform_disable_azure_managed_lustre,
     perform_enable_azure_managed_lustre,
@@ -5382,6 +5385,7 @@ def set_up_azure_monitor_profile(self, mc: ManagedCluster) -> ManagedCluster:
     def set_up_azure_managed_lustre(self, mc: ManagedCluster) -> ManagedCluster:
         self._ensure_mc(mc)
         if self.context.raw_param.get("enable_azure_managed_lustre"):
+            validate_azure_managed_lustre_node_compatibility(self.cmd, mc)
             get_azure_managed_lustre_extension_client(self.cmd)
             self.context.set_intermediate("enable_azure_managed_lustre", True, overwrite_exists=True)
         return mc
@@ -6292,6 +6296,11 @@ def postprocessing_after_mc_created(self, cluster: ManagedCluster) -> None:
                 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()
+            )
+
         # enable azure container storage
         enable_azure_container_storage = self.context.get_intermediate("enable_azure_container_storage")
         container_storage_version = self.context.get_intermediate("container_storage_version")
@@ -6379,11 +6388,6 @@ 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:
@@ -6824,6 +6828,11 @@ def update_azure_managed_lustre(self, mc: ManagedCluster) -> ManagedCluster:
                 self.cmd, self.context.get_resource_group_name(), self.context.get_name()
             )
             validate_azure_managed_lustre_params(enable, disable, installed)
+            if enable and installed:
+                logger.warning("Azure Managed Lustre is already enabled. The existing configuration was preserved.")
+                enable = False
+            elif enable:
+                validate_azure_managed_lustre_node_compatibility(self.cmd, mc)
             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

src/aks-preview/azext_aks_preview/tests/latest/test_azure_managed_lustre.py

@@ -6,7 +6,7 @@
 import inspect
 import unittest
 from types import SimpleNamespace
-from unittest.mock import Mock, patch
+from unittest.mock import Mock, call, patch
 
 from azure.core.exceptions import HttpResponseError
 from azure.cli.core.parser import AzCliCommandParser
@@ -22,7 +22,10 @@
     check_if_extension_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._validators import (
+    validate_azure_managed_lustre_node_compatibility,
+    validate_azure_managed_lustre_params,
+)
 from azext_aks_preview.azuremanagedlustre.aml_ops import (
     perform_disable_azure_managed_lustre,
     perform_enable_azure_managed_lustre,
@@ -107,6 +110,53 @@ def test_valid_parameters(self):
             with self.subTest(enable=enable, disable=disable, installed=installed):
                 validate_azure_managed_lustre_params(enable, disable, installed)
 
+    def test_supported_node_pool(self):
+        cluster = SimpleNamespace(
+            location="eastus",
+            agent_pool_profiles=[SimpleNamespace(os_sku="Ubuntu2404", vm_size="Standard_D2s_v5")],
+        )
+        validate_azure_managed_lustre_node_compatibility(Mock(), cluster)
+
+    def test_unsupported_node_pool_os_skus(self):
+        for os_sku in ["Flatcar", "Ubuntu2604"]:
+            with self.subTest(os_sku=os_sku):
+                cluster = SimpleNamespace(
+                    location="eastus",
+                    agent_pool_profiles=[SimpleNamespace(os_sku=os_sku, vm_size="Standard_D2s_v5")],
+                )
+                with self.assertRaisesRegex(InvalidArgumentValueError, "none of the cluster's node pools"):
+                    validate_azure_managed_lustre_node_compatibility(Mock(), cluster)
+
+    def test_azure_linux_3_node_pool_architecture(self):
+        cluster = SimpleNamespace(
+            location="eastus",
+            agent_pool_profiles=[SimpleNamespace(os_sku="AzureLinux3", vm_size="Standard_D2ps_v5")],
+        )
+        with patch(
+            "azext_aks_preview.azuremanagedlustre._validators._get_vm_sku_architectures",
+            return_value={"standard_d2ps_v5": "arm64"},
+        ):
+            with self.assertRaisesRegex(InvalidArgumentValueError, "ARM64"):
+                validate_azure_managed_lustre_node_compatibility(Mock(), cluster)
+
+        with patch(
+            "azext_aks_preview.azuremanagedlustre._validators._get_vm_sku_architectures",
+            return_value={"standard_d2ps_v5": "x64"},
+        ):
+            validate_azure_managed_lustre_node_compatibility(Mock(), cluster)
+
+    def test_unknown_azure_linux_3_node_pool_architecture(self):
+        cluster = SimpleNamespace(
+            location="eastus",
+            agent_pool_profiles=[SimpleNamespace(os_sku="AzureLinux3", vm_size="Standard_Unknown")],
+        )
+        with patch(
+            "azext_aks_preview.azuremanagedlustre._validators._get_vm_sku_architectures",
+            return_value={},
+        ):
+            with self.assertRaisesRegex(InvalidArgumentValueError, "architecture couldn't be determined"):
+                validate_azure_managed_lustre_node_compatibility(Mock(), cluster)
+
 
 class AzureManagedLustreOperationsTestCase(unittest.TestCase):
     def setUp(self):
@@ -186,6 +236,11 @@ def make_decorator(self, create, **params):
         raw_params.update(params)
         decorator = decorator_type(self.cmd, self.client, raw_params, CUSTOM_MGMT_AKS_PREVIEW)
         decorator.context.attach_mc(self.models.ManagedCluster(location="eastus"))
+        decorator.context.mc.agent_pool_profiles = [
+            self.models.ManagedClusterAgentPoolProfile(
+                name="nodepool1", os_sku="Ubuntu2404", vm_size="Standard_D2s_v5"
+            )
+        ]
         decorator.context.set_intermediate("subscription_id", "test-subscription")
         return decorator
 
@@ -282,7 +337,7 @@ def test_update_propagates_lookup_errors_before_cluster_request(self):
         self.client.begin_create_or_update.assert_not_called()
 
     def test_update_prepares_enable_and_disable(self):
-        for enable, installed in [(True, False), (True, True), (False, True)]:
+        for enable, installed in [(True, False), (False, True)]:
             with self.subTest(enable=enable, installed=installed):
                 decorator = self.make_decorator(
                     False, enable_azure_managed_lustre=enable, disable_azure_managed_lustre=not enable
@@ -296,6 +351,16 @@ def test_update_prepares_enable_and_disable(self):
                 self.assertEqual(decorator.context.get_intermediate("disable_azure_managed_lustre"), not enable)
                 self.assertTrue(decorator.check_is_postprocessing_required(mc))
 
+    def test_update_preserves_existing_extension_configuration(self):
+        decorator = self.make_decorator(False, enable_azure_managed_lustre=True)
+        mc = decorator.context.mc
+        with patch(DECORATOR + ".check_if_azure_managed_lustre_is_installed", return_value=True), \
+                patch(DECORATOR + ".validate_azure_managed_lustre_node_compatibility") as validate:
+            self.assertIs(decorator.update_azure_managed_lustre(mc), mc)
+        validate.assert_not_called()
+        self.assertFalse(decorator.context.get_intermediate("enable_azure_managed_lustre"))
+        self.assertFalse(decorator.check_is_postprocessing_required(mc))
+
     def test_postprocessing_performs_requested_operation(self):
         for create, enable in [(True, True), (False, True), (False, False)]:
             with self.subTest(create=create, enable=enable):
@@ -311,6 +376,18 @@ def test_postprocessing_performs_requested_operation(self):
                 (install if enable else uninstall).assert_called_once_with(self.cmd, "rg", "cluster")
                 (uninstall if enable else install).assert_not_called()
 
+    def test_create_installs_lustre_before_container_storage_early_return(self):
+        decorator = self.make_decorator(True)
+        decorator.context.set_intermediate("enable_azure_managed_lustre", True)
+        decorator.context.set_intermediate("enable_azure_container_storage", "ephemeralDisk")
+        decorator.context.set_intermediate("container_storage_version", "1")
+        functions = decorator.context.external_functions
+        with patch.object(functions, "perform_enable_azure_managed_lustre") as install, \
+                patch("azure.cli.command_modules.acs.managed_cluster_decorator."
+                      "AKSManagedClusterUpdateDecorator.postprocessing_after_mc_created"):
+            decorator.postprocessing_after_mc_created(decorator.context.mc)
+        install.assert_called_once_with(self.cmd, "rg", "cluster")
+
     def test_put_waits_for_cluster_before_lustre_even_with_no_wait(self):
         for create in [True, False]:
             with self.subTest(create=create):