[AKS] az aks create/update: Add distributedcache storage option to --enable/--disable-azure-container-storage

Author: rnhan

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 `distributedcache` storage option to `--enable-azure-container-storage` and `--disable-azure-container-storage` to install/uninstall the Azure Container Storage distributed cache controller.
 
 21.0.0b10
 ++++++++

src/aks-preview/azext_aks_preview/_params.py

@@ -267,6 +267,7 @@
     CONST_STORAGE_POOL_TYPE_AZURE_DISK,
     CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK,
     CONST_STORAGE_POOL_TYPE_ELASTIC_SAN,
+    CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE,
     CONST_STORAGE_POOL_SKU_PREMIUM_LRS,
     CONST_STORAGE_POOL_SKU_STANDARD_LRS,
     CONST_STORAGE_POOL_SKU_STANDARDSSD_LRS,
@@ -480,12 +481,14 @@
     CONST_STORAGE_POOL_TYPE_AZURE_DISK,
     CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK,
     CONST_STORAGE_POOL_TYPE_ELASTIC_SAN,
+    CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE,
 ]
 
 disable_storage_pool_types = [
     CONST_STORAGE_POOL_TYPE_AZURE_DISK,
     CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK,
     CONST_STORAGE_POOL_TYPE_ELASTIC_SAN,
+    CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE,
     CONST_ACSTOR_ALL,
 ]
 
@@ -1198,7 +1201,7 @@ def load_arguments(self, _):
             "enable_azure_container_storage",
             arg_type=_get_container_storage_enum_type(storage_pool_types),
             help="enable azure container storage. Can be used as a flag (defaults to True) or with a"
-            " storage pool type value: (azureDisk, ephemeralDisk, elasticSan)",
+            " storage pool type value: (azureDisk, ephemeralDisk, elasticSan, distributedcache)",
         )
         c.argument(
             "container_storage_version",
@@ -1907,14 +1910,14 @@ def load_arguments(self, _):
             "enable_azure_container_storage",
             arg_type=_get_container_storage_enum_type(storage_pool_types),
             help="enable azure container storage. Can be used as a flag (defaults to True) or with a"
-            " storage pool type value: (azureDisk, ephemeralDisk, elasticSan)",
+            " storage pool type value: (azureDisk, ephemeralDisk, elasticSan, distributedcache)",
         )
         c.argument(
             "disable_azure_container_storage",
             arg_type=_get_container_storage_enum_type(disable_storage_pool_types),
-            help="disable azure container storage or any one of the storage pool types."
-            " Can be used as a flag (defaults to True) or with a storagepool type value:"
-            " azureDisk, ephemeralDisk, elasticSan, all (to disable all storage pools).",
+            help="disable azure container storage or any one of the storage types."
+            " Can be used as a flag (defaults to True) or with a storage type value:"
+            " azureDisk, ephemeralDisk, elasticSan, distributedcache, all (to disable all storage types).",
         )
         c.argument(
             "container_storage_version",

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

@@ -17,6 +17,10 @@
 CONST_ACSTOR_EXT_INSTALLATION_NAME = "acstor"
 CONST_ACSTOR_EXT_INSTALLATION_NAMESPACE = "kube-system"
 CONST_ACSTOR_K8S_EXTENSION_NAME = "microsoft.azurecontainerstoragev2"
+# Distributed cache (Distributed Accelerator for Cloud Storage) install controller.
+CONST_DISTRIBUTED_CACHE_EXT_INSTALLATION_NAME = "distributedcache"
+CONST_DISTRIBUTED_CACHE_EXT_INSTALLATION_NAMESPACE = "kube-system"
+CONST_DISTRIBUTED_CACHE_K8S_EXTENSION_NAME = "microsoft.dacs"
 CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME = "azext_k8s_extension._client_factory"
 CONST_K8S_EXTENSION_CUSTOM_MOD_NAME = "azext_k8s_extension.custom"
 CONST_K8S_EXTENSION_NAME = "k8s-extension"
@@ -35,3 +39,4 @@
 CONST_STORAGE_POOL_TYPE_AZURE_DISK = "azureDisk"
 CONST_STORAGE_POOL_TYPE_ELASTIC_SAN = "elasticSan"
 CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK = "ephemeralDisk"
+CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE = "distributedcache"

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

@@ -23,6 +23,7 @@
     CONST_STORAGE_POOL_TYPE_AZURE_DISK,
     CONST_STORAGE_POOL_TYPE_ELASTIC_SAN,
     CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK,
+    CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE,
     CONST_ACSTOR_V1_K8S_EXTENSION_NAME,
     CONST_ACSTOR_V1_EXT_INSTALLATION_NAME,
 )
@@ -304,6 +305,15 @@ def should_delete_extension(storage_options_to_remove) -> bool:
     )
 
 
+def is_distributed_cache_requested(storage_options) -> bool:
+    # Detect whether distributed cache is among the requested storage options.
+    if storage_options is None or storage_options is True:
+        return False
+    if isinstance(storage_options, list):
+        return CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE in storage_options
+    return storage_options == CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE
+
+
 def get_container_storage_extension_installed(
     cmd,
     resource_group,

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

@@ -21,6 +21,7 @@
     CONST_STORAGE_POOL_TYPE_AZURE_DISK,
     CONST_STORAGE_POOL_TYPE_ELASTIC_SAN,
     CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK,
+    CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE,
 )
 from azext_aks_preview.azurecontainerstorage._helpers import (
     get_vm_sku_details
@@ -571,6 +572,114 @@ def validate_disable_azure_container_storage_params(
         )
 
 
+def validate_enable_distributed_cache_params(
+    enablement_option,
+    is_extension_installed,
+    storage_pool_name,
+    storage_pool_sku,
+    storage_pool_option,
+    storage_pool_size,
+    ephemeral_disk_volume_type,
+    ephemeral_disk_nvme_perf_tier,
+    container_storage_version=None,
+):
+    # Distributed cache has no storage pool construct, so none of the storage
+    # pool parameters are supported.
+    enablement_option_arr = enablement_option if isinstance(enablement_option, list) else [enablement_option]
+    other_pool_types = [
+        opt for opt in enablement_option_arr
+        if opt != CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE
+    ]
+    if other_pool_types:
+        options_display = "', '".join(other_pool_types)
+        raise InvalidArgumentValueError(
+            f"'{CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE}' cannot be combined with other storage "
+            f"options ('{options_display}'). Distributed cache is enabled independently. "
+            f"Please run --enable-azure-container-storage {CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE} "
+            "on its own."
+        )
+
+    if is_extension_installed:
+        raise InvalidArgumentValueError(
+            'Cannot enable distributed cache as it is already enabled on the cluster.'
+        )
+
+    unsupported_params = []
+    if storage_pool_name is not None:
+        unsupported_params.append('--storage-pool-name')
+    if storage_pool_sku is not None:
+        unsupported_params.append('--storage-pool-sku')
+    if storage_pool_option is not None:
+        unsupported_params.append('--storage-pool-option')
+    if storage_pool_size is not None:
+        unsupported_params.append('--storage-pool-size')
+    if ephemeral_disk_volume_type is not None:
+        unsupported_params.append('--ephemeral-disk-volume-type')
+    if ephemeral_disk_nvme_perf_tier is not None:
+        unsupported_params.append('--ephemeral-disk-nvme-perf-tier')
+    if container_storage_version is not None:
+        unsupported_params.append('--container-storage-version')
+
+    if unsupported_params:
+        params_defined = ', '.join(unsupported_params)
+        raise InvalidArgumentValueError(
+            f'{params_defined} cannot be used with '
+            f'--enable-azure-container-storage {CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE}. '
+            'Distributed cache does not require or support any storage pool configuration. '
+            'Please remove these parameters and try again.'
+        )
+
+
+def validate_disable_distributed_cache_params(
+    disablement_option,
+    is_extension_installed,
+    storage_pool_name,
+    storage_pool_sku,
+    storage_pool_option,
+    storage_pool_size,
+    container_storage_version=None,
+):
+    disablement_option_arr = disablement_option if isinstance(disablement_option, list) else [disablement_option]
+    other_pool_types = [
+        opt for opt in disablement_option_arr
+        if opt != CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE
+    ]
+    if other_pool_types:
+        options_display = "', '".join(other_pool_types)
+        raise InvalidArgumentValueError(
+            f"'{CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE}' cannot be combined with other storage "
+            f"options ('{options_display}'). Distributed cache is disabled independently. "
+            f"Please run --disable-azure-container-storage {CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE} "
+            "on its own."
+        )
+
+    if not is_extension_installed:
+        raise InvalidArgumentValueError(
+            'Cannot disable distributed cache as it could not be found on the cluster.'
+        )
+
+    unsupported_params = []
+    if storage_pool_name is not None:
+        unsupported_params.append('--storage-pool-name')
+    if storage_pool_sku is not None:
+        unsupported_params.append('--storage-pool-sku')
+    if storage_pool_option is not None:
+        unsupported_params.append('--storage-pool-option')
+    if storage_pool_size is not None:
+        unsupported_params.append('--storage-pool-size')
+    if container_storage_version is not None:
+        unsupported_params.append('--container-storage-version')
+
+    if unsupported_params:
+        params_defined = ', '.join(unsupported_params)
+        raise InvalidArgumentValueError(
+            f'{params_defined} cannot be used with '
+            f'--disable-azure-container-storage {CONST_STORAGE_POOL_TYPE_DISTRIBUTED_CACHE}. '
+            'Distributed cache does not require or support any storage pool configuration. '
+            'Please remove these parameters and try again.'
+        )
+
+
 # _Validate_storage_pool_size validates that the storage_pool_size is
 # string of a combination of a float number immediately followed by
 # Ti or Gi e.g. 2Ti, 512Gi, 1.5Ti. The function also validates that the

src/aks-preview/azext_aks_preview/azurecontainerstorage/acstor_ops.py

@@ -12,6 +12,9 @@
     CONST_ACSTOR_EXT_INSTALLATION_NAME,
     CONST_ACSTOR_EXT_INSTALLATION_NAMESPACE,
     CONST_ACSTOR_K8S_EXTENSION_NAME,
+    CONST_DISTRIBUTED_CACHE_EXT_INSTALLATION_NAME,
+    CONST_DISTRIBUTED_CACHE_EXT_INSTALLATION_NAMESPACE,
+    CONST_DISTRIBUTED_CACHE_K8S_EXTENSION_NAME,
     CONST_DISK_TYPE_EPHEMERAL_VOLUME_ONLY,
     CONST_DISK_TYPE_PV_WITH_ANNOTATION,
     CONST_EPHEMERAL_NVME_PERF_TIER_STANDARD,
@@ -797,3 +800,109 @@ def perform_azure_container_storage_update(
             raise UnknownError(
                 "Failed to disable Azure Container Storage with error: %s" % delete_ex
             ) from delete_ex
+
+
+def perform_enable_distributed_cache(
+    cmd,
+    resource_group,
+    cluster_name,
+    is_extension_installed=False,
+    is_called_from_extension=False,
+):
+    # This will be set true only when aks-preview extension is used
+    # and we want the aks-preview ManagedClusterDecorator to call the
+    # perform_enable_distributed_cache function.
+    if not is_called_from_extension:
+        return
+
+    # Enabling only installs the install controller. Cache components are
+    # deployed later, driven by the creation of the corresponding CR.
+    client_factory = get_k8s_extension_module(CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME)
+    client = client_factory.cf_k8s_extension_operation(cmd.cli_ctx)
+    k8s_extension_custom_mod = get_k8s_extension_module(CONST_K8S_EXTENSION_CUSTOM_MOD_NAME)
+
+    if is_extension_installed:
+        logger.warning("Distributed cache is already enabled on the cluster.")
+        return
+
+    try:
+        result = k8s_extension_custom_mod.create_k8s_extension(
+            cmd,
+            client,
+            resource_group,
+            cluster_name,
+            CONST_DISTRIBUTED_CACHE_EXT_INSTALLATION_NAME,
+            "managedClusters",
+            CONST_DISTRIBUTED_CACHE_K8S_EXTENSION_NAME,
+            auto_upgrade_minor_version=True,
+            release_train="stable",
+            scope="cluster",
+            release_namespace=CONST_DISTRIBUTED_CACHE_EXT_INSTALLATION_NAMESPACE,
+        )
+        long_op_result = LongRunningOperation(cmd.cli_ctx)(result)
+        if long_op_result.provisioning_state == "Succeeded":
+            logger.warning("Distributed cache successfully installed")
+    except Exception as ex:  # pylint: disable=broad-except
+        logger.error("Distributed cache failed to install.\nError: %s", ex)
+        logger.warning("Cleaning up the cluster by disabling distributed cache")
+        try:
+            delete_op_result = k8s_extension_custom_mod.delete_k8s_extension(
+                cmd,
+                client,
+                resource_group,
+                cluster_name,
+                CONST_DISTRIBUTED_CACHE_EXT_INSTALLATION_NAME,
+                "managedClusters",
+                yes=True,
+            )
+            LongRunningOperation(cmd.cli_ctx)(delete_op_result)
+            logger.warning(
+                "Please retry enabling distributed cache by running "
+                "`az aks update` along with "
+                "`--enable-azure-container-storage distributedcache`"
+            )
+        except Exception as delete_ex:  # pylint: disable=broad-except
+            raise UnknownError(
+                "Failed to clean up distributed cache with error: %s" % delete_ex
+            ) from delete_ex
+
+
+def perform_disable_distributed_cache(
+    cmd,
+    resource_group,
+    cluster_name,
+    is_extension_installed=False,
+    is_called_from_extension=False,
+):
+    # This will be set true only when aks-preview extension is used
+    # and we want the aks-preview ManagedClusterDecorator to call the
+    # perform_disable_distributed_cache function.
+    if not is_called_from_extension:
+        return
+
+    client_factory = get_k8s_extension_module(CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME)
+    client = client_factory.cf_k8s_extension_operation(cmd.cli_ctx)
+    k8s_extension_custom_mod = get_k8s_extension_module(CONST_K8S_EXTENSION_CUSTOM_MOD_NAME)
+
+    if not is_extension_installed:
+        logger.warning("Distributed cache is not enabled on the cluster.")
+        return
+
+    # Deleting the install controller triggers removal of the remaining cache
+    # components from the cluster.
+    try:
+        delete_op_result = k8s_extension_custom_mod.delete_k8s_extension(
+            cmd,
+            client,
+            resource_group,
+            cluster_name,
+            CONST_DISTRIBUTED_CACHE_EXT_INSTALLATION_NAME,
+            "managedClusters",
+            yes=True,
+        )
+        LongRunningOperation(cmd.cli_ctx)(delete_op_result)
+        logger.warning("Distributed cache has been disabled.")
+    except Exception as delete_ex:  # pylint: disable=broad-except
+        raise UnknownError(
+            "Failed to disable distributed cache with error: %s" % delete_ex
+        ) from delete_ex