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

src/aks-preview/HISTORY.rst

@@ -11,6 +11,10 @@ To release a new version, please select a new version number (usually plus 1 to
 
 Pending
 +++++++
+
+22.0.0b9
++++++++++
+* `az aks create/update`: Add `distributedAccelerator` storage option to `--enable-azure-container-storage` and `--disable-azure-container-storage` to install/uninstall the Azure Container Storage distributed accelerator controller.
 * `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/azext_aks_preview/_params.py

@@ -280,6 +280,7 @@
     CONST_STORAGE_POOL_TYPE_AZURE_DISK,
     CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK,
     CONST_STORAGE_POOL_TYPE_ELASTIC_SAN,
+    CONST_STORAGE_POOL_TYPE_DISTRIBUTED_ACCELERATOR,
     CONST_STORAGE_POOL_SKU_PREMIUM_LRS,
     CONST_STORAGE_POOL_SKU_STANDARD_LRS,
     CONST_STORAGE_POOL_SKU_STANDARDSSD_LRS,
@@ -529,12 +530,14 @@ def _deprecate_option(c, target, redirect):
     CONST_STORAGE_POOL_TYPE_AZURE_DISK,
     CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK,
     CONST_STORAGE_POOL_TYPE_ELASTIC_SAN,
+    CONST_STORAGE_POOL_TYPE_DISTRIBUTED_ACCELERATOR,
 ]
 
 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_ACCELERATOR,
     CONST_ACSTOR_ALL,
 ]
 
@@ -1294,7 +1297,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, distributedAccelerator)",
         )
         c.argument(
             "container_storage_version",
@@ -2042,14 +2045,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, distributedAccelerator)",
         )
         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, distributedAccelerator, 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 accelerator (Distributed Accelerator for Cloud Storage) install controller.
+CONST_DISTRIBUTED_ACCELERATOR_EXT_INSTALLATION_NAME = "distributedaccelerator"
+CONST_DISTRIBUTED_ACCELERATOR_EXT_INSTALLATION_NAMESPACE = "kube-system"
+CONST_DISTRIBUTED_ACCELERATOR_K8S_EXTENSION_NAME = "microsoft.distributedaccelerator"
 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_ACCELERATOR = "distributedAccelerator"

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_ACCELERATOR,
     CONST_ACSTOR_V1_K8S_EXTENSION_NAME,
     CONST_ACSTOR_V1_EXT_INSTALLATION_NAME,
 )
@@ -304,11 +305,21 @@ def should_delete_extension(storage_options_to_remove) -> bool:
     )
 
 
+def is_distributed_accelerator_requested(storage_options) -> bool:
+    # Detect whether distributed accelerator 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_ACCELERATOR in storage_options
+    return storage_options == CONST_STORAGE_POOL_TYPE_DISTRIBUTED_ACCELERATOR
+
+
 def get_container_storage_extension_installed(
     cmd,
     resource_group,
     cluster_name,
     extension_name,
+    expected_extension_type=None,
 ) -> Tuple[bool, str]:
 
     client_factory = get_k8s_extension_module(CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME)
@@ -325,6 +336,13 @@ def get_container_storage_extension_installed(
             extension_name,
             "managedClusters",
         )
+        # show_k8s_extension only matches on the instance name, so an unrelated
+        # extension that happens to share the name would otherwise be treated as
+        # installed. When an expected extension type is supplied, verify it before
+        # reporting the extension as installed to avoid acting on the wrong resource.
+        if expected_extension_type is not None and \
+                (getattr(extension, "extension_type", None) or "").lower() != expected_extension_type.lower():
+            return False, ""
         is_extension_installed = True
         extension_version = extension.current_version
     except ResourceNotFoundError:

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_ACCELERATOR,
 )
 from azext_aks_preview.azurecontainerstorage._helpers import (
     get_vm_sku_details
@@ -571,6 +572,163 @@ def validate_disable_azure_container_storage_params(
         )
 
 
+def _reject_unsupported_distributed_accelerator_params(
+    command_display,
+    storage_pool_name,
+    storage_pool_sku,
+    storage_pool_option,
+    storage_pool_size,
+    ephemeral_disk_volume_type,
+    ephemeral_disk_nvme_perf_tier,
+    nodepool_list,
+    container_storage_version,
+):
+    # Distributed accelerator has no storage pool construct, so none of the storage
+    # pool scoping parameters are supported on the distributed accelerator paths.
+    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 nodepool_list is not None:
+        unsupported_params.append('--azure-container-storage-nodepools')
+    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 {command_display}. '
+            'Distributed accelerator does not require or support any storage pool configuration. '
+            'Please remove these parameters and try again.'
+        )
+
+
+def validate_enable_distributed_accelerator_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,
+    nodepool_list=None,
+    container_storage_version=None,
+):
+    # Distributed accelerator 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_ACCELERATOR
+    ]
+    if other_pool_types:
+        options_display = "', '".join(other_pool_types)
+        raise InvalidArgumentValueError(
+            f"'{CONST_STORAGE_POOL_TYPE_DISTRIBUTED_ACCELERATOR}' cannot be combined with other storage "
+            f"options ('{options_display}'). Distributed accelerator is enabled independently. "
+            f"Please run --enable-azure-container-storage {CONST_STORAGE_POOL_TYPE_DISTRIBUTED_ACCELERATOR} "
+            "on its own."
+        )
+
+    if is_extension_installed:
+        raise InvalidArgumentValueError(
+            'Cannot enable distributed accelerator as it is already enabled on the cluster.'
+        )
+
+    _reject_unsupported_distributed_accelerator_params(
+        f'--enable-azure-container-storage {CONST_STORAGE_POOL_TYPE_DISTRIBUTED_ACCELERATOR}',
+        storage_pool_name,
+        storage_pool_sku,
+        storage_pool_option,
+        storage_pool_size,
+        ephemeral_disk_volume_type,
+        ephemeral_disk_nvme_perf_tier,
+        nodepool_list,
+        container_storage_version,
+    )
+
+
+def validate_disable_distributed_accelerator_params(
+    disablement_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,
+    nodepool_list=None,
+    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_ACCELERATOR
+    ]
+    if other_pool_types:
+        options_display = "', '".join(other_pool_types)
+        raise InvalidArgumentValueError(
+            f"'{CONST_STORAGE_POOL_TYPE_DISTRIBUTED_ACCELERATOR}' cannot be combined with other storage "
+            f"options ('{options_display}'). Distributed accelerator is disabled independently. "
+            f"Please run --disable-azure-container-storage {CONST_STORAGE_POOL_TYPE_DISTRIBUTED_ACCELERATOR} "
+            "on its own."
+        )
+
+    if not is_extension_installed:
+        raise InvalidArgumentValueError(
+            'Cannot disable distributed accelerator as it could not be found on the cluster.'
+        )
+
+    _reject_unsupported_distributed_accelerator_params(
+        f'--disable-azure-container-storage {CONST_STORAGE_POOL_TYPE_DISTRIBUTED_ACCELERATOR}',
+        storage_pool_name,
+        storage_pool_sku,
+        storage_pool_option,
+        storage_pool_size,
+        ephemeral_disk_volume_type,
+        ephemeral_disk_nvme_perf_tier,
+        nodepool_list,
+        container_storage_version,
+    )
+
+
+def validate_disable_all_distributed_accelerator_params(
+    storage_pool_name,
+    storage_pool_sku,
+    storage_pool_option,
+    storage_pool_size,
+    ephemeral_disk_volume_type,
+    ephemeral_disk_nvme_perf_tier,
+    nodepool_list=None,
+    container_storage_version=None,
+):
+    # A bare `--disable-azure-container-storage` or `all` teardown of a
+    # distributed-accelerator-only cluster removes the controller directly and
+    # returns before the storage pool validation runs, so reject the storage
+    # pool scoping parameters here as well.
+    _reject_unsupported_distributed_accelerator_params(
+        '--disable-azure-container-storage on a distributed accelerator cluster',
+        storage_pool_name,
+        storage_pool_sku,
+        storage_pool_option,
+        storage_pool_size,
+        ephemeral_disk_volume_type,
+        ephemeral_disk_nvme_perf_tier,
+        nodepool_list,
+        container_storage_version,
+    )
+
+
 # _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