diff --git a/src/aks-preview/azext_aks_preview/azurecontainerstorage/_helpers.py b/src/aks-preview/azext_aks_preview/azurecontainerstorage/_helpers.py index aa2a9a9ccd..e0592e61e5 100644 --- a/src/aks-preview/azext_aks_preview/azurecontainerstorage/_helpers.py +++ b/src/aks-preview/azext_aks_preview/azurecontainerstorage/_helpers.py @@ -30,12 +30,13 @@ delete_role_assignments, ) from azure.cli.core.azclierror import UnknownError -from cachetools import cached -from cachetools.keys import hashkey from knack.log import get_logger logger = get_logger(__name__) +# Global cache to store VM SKU information. +vm_sku_details_cache = {} + def validate_storagepool_creation( cmd, @@ -453,15 +454,11 @@ def generate_vm_sku_cache_for_region(cli_ctx, location=None): if entry.name == 'NvmeDiskSizeInMiB': nvme_enabled = True - vm_sku_details(sku_name, cpu_value, nvme_enabled) + vm_sku_details_cache[sku_name] = (cpu_value, nvme_enabled) -@cached( - cache={}, - key=lambda sku_name, cpu_value=None, nvme_enabled=None: hashkey(sku_name) -) -def vm_sku_details(sku_name, cpu_value=None, nvme_enabled=None): # pylint: disable=unused-argument - return cpu_value, nvme_enabled +def get_vm_sku_details(sku_name): + return vm_sku_details_cache.get(sku_name) def _get_ephemeral_nvme_cpu_value_based_on_vm_size_perf_tier(nodepool_skus, perf_tier): @@ -472,7 +469,11 @@ def _get_ephemeral_nvme_cpu_value_based_on_vm_size_perf_tier(nodepool_skus, perf elif perf_tier.lower() == CONST_EPHEMERAL_NVME_PERF_TIER_PREMIUM.lower(): multiplication_factor = 0.5 for vm_size in nodepool_skus: - number_of_cores, _ = vm_sku_details(vm_size.lower()) + number_of_cores, _ = get_vm_sku_details(vm_size.lower()) + if number_of_cores is None: + raise UnknownError( + f"Unable to find details for virtual machine size {vm_size}." + ) if number_of_cores != -1: if cpu_value == -1: cpu_value = number_of_cores * multiplication_factor diff --git a/src/aks-preview/azext_aks_preview/azurecontainerstorage/_validators.py b/src/aks-preview/azext_aks_preview/azurecontainerstorage/_validators.py index 66a94ba5eb..29e7433b7d 100644 --- a/src/aks-preview/azext_aks_preview/azurecontainerstorage/_validators.py +++ b/src/aks-preview/azext_aks_preview/azurecontainerstorage/_validators.py @@ -22,7 +22,7 @@ CONST_STORAGE_POOL_TYPE_EPHEMERAL_DISK, ) from azext_aks_preview.azurecontainerstorage._helpers import ( - vm_sku_details + get_vm_sku_details ) from azure.cli.core.azclierror import ( ArgumentUsageError, @@ -516,7 +516,11 @@ def _validate_nodepools( # pylint: disable=too-many-branches,too-many-locals ) vm_size = agentpool.get("vm_size") if vm_size is not None: - cpu_value, nvme_enabled = vm_sku_details(vm_size.lower()) + cpu_value, nvme_enabled = get_vm_sku_details(vm_size.lower()) + if cpu_value is None or nvme_enabled is None: + raise UnknownError( + f'Unable to find details for virtual machine size {vm_size}.' + ) if cpu_value < 0: raise UnknownError( f'Unable to determine number of cores in node pool: {pool_name}, node size: {vm_size}' diff --git a/src/aks-preview/azext_aks_preview/tests/latest/test_validators.py b/src/aks-preview/azext_aks_preview/tests/latest/test_validators.py index 9a7cdb6f68..b83b11a7ad 100644 --- a/src/aks-preview/azext_aks_preview/tests/latest/test_validators.py +++ b/src/aks-preview/azext_aks_preview/tests/latest/test_validators.py @@ -875,16 +875,7 @@ def test_valid_disable(self): class TestValidateEnableAzureContainerStorage(unittest.TestCase): @classmethod def setUpClass(cls): - cls.mock_cache = { - ('standard_l8s_v3',): (8, True), - ('standard_d2s_v2',): (2, False), - ('standard_d2pds_v6',): (2, True), - ('standard_ds1_v2',): (1, False), - ('standard_m8-2ms',): (2, False), - ('standard_b2s',): (2, False), - } - - def side_effect_fn(sku_name, cpu_val=None, nvme=None): + def side_effect_fn(sku_name): if sku_name == "standard_l8s_v3": return 8, True elif sku_name == "standard_d2s_v2": @@ -900,7 +891,7 @@ def side_effect_fn(sku_name, cpu_val=None, nvme=None): return None, None - cls.patcher = patch('azext_aks_preview.azurecontainerstorage._validators.vm_sku_details') + cls.patcher = patch('azext_aks_preview.azurecontainerstorage._validators.get_vm_sku_details') cls.mock_fn = cls.patcher.start() cls.mock_fn.side_effect = side_effect_fn