{aks-preview} Migrate commands calling Compute module to aaz-based implementation

src/aks-preview/HISTORY.rst

@@ -4,14 +4,18 @@ Release History
 ===============
 
 Guidance
-++++++++
++++++++
 If there is no rush to release a new version, please just add a description of the modification under the *Pending* section.
 
 To release a new version, please select a new version number (usually plus 1 to last patch version, X.Y.Z -> Major.Minor.Patch, more details in `\doc <https://semver.org/>`_), and then add a new section named as the new version number in this file, the content should include the new modifications and everything from the *Pending* section. Finally, update the `VERSION` variable in `setup.py` with this new version number.
 
 Pending
 +++++++
 
+20.0.0b3
++++++++
+* Migrate code from Azure SDK to AAZ based commands for compute operations (VM).
+
 20.0.0b2
 +++++++
 * `az aks nodepool update`: clean up some useless code in the update managed gpu function. 

src/aks-preview/azext_aks_preview/_client_factory.py

@@ -71,10 +71,6 @@ def get_mc_snapshots_client(cli_ctx, subscription_id=None):
     return get_container_service_client(cli_ctx, subscription_id=subscription_id).managed_cluster_snapshots
 
 
-def get_compute_client(cli_ctx, *_):
-    return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_COMPUTE)
-
-
 def get_resource_groups_client(cli_ctx, subscription_id=None):
     return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES,
                                    subscription_id=subscription_id).resource_groups

src/aks-preview/azext_aks_preview/_completers.py

@@ -53,13 +53,14 @@ def get_vm_size_completion_list(cmd, prefix, namespace, **kwargs):  # pylint: di
 
     location = _get_location(cmd.cli_ctx, namespace)
     result = get_vm_sizes(cmd.cli_ctx, location)
-    return set(r.name for r in result) & set(c.value for c in ContainerServiceVMSizeTypes)
+    return set(r.get('name') for r in result) & set(c.value for c in ContainerServiceVMSizeTypes)
 
 
 def get_vm_sizes(cli_ctx, location):
-    from azext_aks_preview._client_factory import get_compute_client
-
-    return get_compute_client(cli_ctx).virtual_machine_sizes.list(location)
+    from azure.cli.command_modules.vm.operations.vm import VMListSizes
+    return VMListSizes(cli_ctx=cli_ctx)(command_args={
+        "location": location,
+    })
 
 
 def _get_location(cli_ctx, namespace):

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

@@ -528,18 +528,18 @@ def check_if_new_storagepool_creation_required(
 def generate_vm_sku_cache_for_region(cli_ctx, location=None):
     result = _get_vm_sku_details(cli_ctx, location)
     for vm_data in result:
-        sku_name = vm_data.name.lower()
-        capabilities = vm_data.capabilities
+        sku_name = vm_data.get('name', '').lower()
+        capabilities = vm_data.get('capabilities', [])
         cpu_value = -1
         nvme_enabled = False
         for entry in capabilities:
-            if entry.name == 'vCPUs' and cpu_value == -1:
-                cpu_value = int(entry.value)
+            if entry.get('name') == 'vCPUs' and cpu_value == -1:
+                cpu_value = int(entry.get('value'))
 
-            if entry.name == 'vCPUsAvailable':
-                cpu_value = int(entry.value)
+            if entry.get('name') == 'vCPUsAvailable':
+                cpu_value = int(entry.get('value'))
 
-            if entry.name == 'NvmeDiskSizeInMiB':
+            if entry.get('name') == 'NvmeDiskSizeInMiB':
                 nvme_enabled = True
 
             vm_sku_details_cache[sku_name] = (cpu_value, nvme_enabled)
@@ -652,9 +652,9 @@ def _is_vm_in_required_location(desired_location, location_list):
                 return True
         return False
 
-    from azext_aks_preview._client_factory import get_compute_client
-    result = get_compute_client(cli_ctx).resource_skus.list()
-    result = [x for x in result if x.resource_type.lower() == 'virtualmachines']
+    from azure.cli.command_modules.vm.aaz.latest.vm import ListSkus as VMListSkus
+    result = VMListSkus(cli_ctx=cli_ctx)(command_args={})
+    result = [x for x in result if x.get('resourceType', '').lower() == 'virtualmachines']
     if location:
-        result = [r for r in result if _is_vm_in_required_location(location, r.locations)]
+        result = [r for r in result if _is_vm_in_required_location(location, r.get('locations', []))]
     return result

src/aks-preview/azext_aks_preview/custom.py

@@ -17,11 +17,7 @@
 import subprocess
 from math import isnan
 
-from azext_aks_preview._client_factory import (
-    CUSTOM_MGMT_AKS_PREVIEW,
-    cf_agent_pools,
-    get_compute_client,
-)
+from azext_aks_preview._client_factory import CUSTOM_MGMT_AKS_PREVIEW, cf_agent_pools
 from azext_aks_preview._consts import (
     ADDONS,
     ADDONS_DESCRIPTIONS,
@@ -4670,29 +4666,34 @@ def _aks_run_command(
             command += f" {all_endpoints}"
             logger.debug("Full command: %s", command)
 
-        compute_client = get_compute_client(cmd.cli_ctx)
-
         if vm_set_type == CONST_VIRTUAL_MACHINE_SCALE_SETS:
-            RunCommandInput = cmd.get_models('RunCommandInput',
-                                             resource_type=ResourceType.MGMT_COMPUTE,
-                                             operation_group="virtual_machine_scale_sets")
+            from azure.cli.command_modules.vm.aaz.latest.vmss.run_command import Invoke as VMSSRunInvoke
+            command_args = {
+                'instance_id': instance_id,
+                'resource_group': managed_resource_group,
+                'vmss_name': vmss_name,
+                'command_id': "RunShellScript",
+                'script': [command]
+            }
             command_result = LongRunningOperation(cmd.cli_ctx)(
-                compute_client.virtual_machine_scale_set_vms.begin_run_command(
-                    managed_resource_group, vmss_name, instance_id,
-                    RunCommandInput(command_id="RunShellScript", script=[command])))
+                VMSSRunInvoke(cli_ctx=cmd.cli_ctx)(command_args=command_args))
         elif vm_set_type == CONST_AVAILABILITY_SET:
-            RunCommandInput = cmd.get_models('RunCommandInput',
-                                             resource_type=ResourceType.MGMT_COMPUTE,
-                                             operation_group="virtual_machine_run_commands")
+            from azure.cli.command_modules.vm.aaz.latest.vm.run_command import Invoke as VMRunInvoke
+            command_args = {
+                'resource_group': managed_resource_group,
+                'vm_name': vm_name,
+                'command_id': "RunShellScript",
+                'script': [command]
+            }
             command_result = LongRunningOperation(cmd.cli_ctx)(
-                compute_client.virtual_machines.begin_run_command(
-                    managed_resource_group, vm_name,
-                    RunCommandInput(command_id="RunShellScript", script=[command])))
+                VMRunInvoke(cli_ctx=cmd.cli_ctx)(command_args=command_args))
         else:
             raise ValidationError(f"VM set type {vm_set_type} is not supported!")
 
-        display_status = command_result.value[0].display_status
-        message = command_result.value[0].message
+        display_status, message = '', ''
+        if len(command_result.get('value', [])) > 0:
+            display_status = command_result['value'][0].get('displayStatus')
+            message = command_result['value'][0].get('message')
         if display_status != "Provisioning succeeded":
             raise InvalidArgumentValueError(
                 f"Can not run command with returned code {display_status} and message {message}")
@@ -4730,13 +4731,44 @@ def _aks_verify_resource(resource, resource_type):
                               f"Image version must be at least {CONST_MIN_NODE_IMAGE_VERSION}.")
 
 
+def _aks_verify_resource_by_aaz(resource, resource_type):
+    if resource.get('provisioningState') != CONST_NODE_PROVISIONING_STATE_SUCCEEDED:
+        raise ValidationError(f"Node pool {resource.get('name')} is in {resource.get('provisioningState')} state!")
+
+    node_image_version = ""
+    os_type = ""
+    if resource_type == CONST_VIRTUAL_MACHINE_SCALE_SETS:
+        node_image_version = resource.get('nodeImageVersion')
+        os_type = resource.get('osType')
+    else:
+        node_image_version = resource.get('storageProfile', {}).get('imageReference', {}).get('id')
+        os_type = resource.get('storageProfile', {}).get('osDisk', {}).get('osType')
+
+    if not os_type or os_type != CONST_DEFAULT_NODE_OS_TYPE:
+        raise ValidationError(f"Resource must be of type {CONST_DEFAULT_NODE_OS_TYPE}!")
+
+    if not node_image_version:
+        raise ValidationError(f"No image version found for {resource.get('name')}! Cannot verify supported versions.")
+
+    if resource_type == CONST_VIRTUAL_MACHINE_SCALE_SETS:
+        version = node_image_version.split("-")[-1]
+    else:
+        version = node_image_version.split("/")[-1]
+
+    if version < CONST_MIN_NODE_IMAGE_VERSION:
+        raise ValidationError(f"Node image version {version} is not supported! "
+                              f"Image version must be at least {CONST_MIN_NODE_IMAGE_VERSION}.")
+
+
 def _aks_get_node_name_vmss(
         cmd,
         resource_group,
         cluster_name,
         node_name,
         managed_resource_group):
-    compute_client = get_compute_client(cmd.cli_ctx)
+    from azure.cli.command_modules.vm.aaz.latest.vmss import List as VMSSList
+    from azure.cli.command_modules.vm.operations.vmss import VMSSListInstances
+    from azure.cli.command_modules.vm.operations.vmss_vms import VMSSVMSShow
 
     if not node_name:
         logger.debug("No node name specified, will randomly select a node from the cluster")
@@ -4760,35 +4792,45 @@ def _aks_get_node_name_vmss(
         if not nodepool_name:
             raise ValidationError("No suitable node pool found in the cluster.")
 
-        vmss_list = compute_client.virtual_machine_scale_sets.list(managed_resource_group)
+        vmss_list = VMSSList(cli_ctx=cmd.cli_ctx)({
+            'resource_group': managed_resource_group
+        })
         if not vmss_list:
             raise ValidationError(f"No VMSS found in the managed resource group {managed_resource_group}!")
 
         for vmss in vmss_list:
-            vmss_tag = vmss.tags.get("aks-managed-poolName")
+            vmss_tag = vmss.get('tags', {}).get("aks-managed-poolName")
             if vmss_tag and vmss_tag == nodepool_name:
-                vmss_name = vmss.name
+                vmss_name = vmss.get('name')
                 logger.debug("Select VMSS: %s", vmss_name)
                 break
         if not vmss_name:
             raise ValidationError(f"No VMSS pool matched AKS node pool {nodepool_name}!")
 
-        instances = list(compute_client.virtual_machine_scale_set_vms.list(managed_resource_group, vmss_name))
-        if not instances:
+        command_args = {
+            'resource_group': managed_resource_group,
+            'virtual_machine_scale_set_name': vmss_name
+        }
+        instances = list(VMSSListInstances(cli_ctx=cmd.cli_ctx)(command_args=command_args))
+        if not instances or len(instances) < 1:
             raise ValidationError(f"No instances found in the VMSS {vmss_name}!")
 
-        instance_id = instances[0].instance_id
+        instance_id = instances[0].get('instanceId')
         logger.debug("Select instance id: %s", instance_id)
     else:
         index = node_name.find("vmss")
         if index != -1:
             vmss_name = node_name[:index + 4]
-            instance_id = int(node_name[index + 4:], base=36)
-            instance_info = compute_client.virtual_machine_scale_set_vms.get(
-                managed_resource_group, vmss_name, instance_id)
+            instance_id = str(int(node_name[index + 4:], base=36))
+            command_args = {
+                'instance_id': instance_id,
+                'resource_group': managed_resource_group,
+                'vm_scale_set_name': vmss_name
+            }
+            instance_info = VMSSVMSShow(cli_ctx=cmd.cli_ctx)(command_args=command_args)
             if not instance_info:
                 raise ValidationError(f"Instance id {instance_id} not found in VMSS {vmss_name}!")
-            _aks_verify_resource(instance_info, CONST_VIRTUAL_MACHINES)
+            _aks_verify_resource_by_aaz(instance_info, CONST_VIRTUAL_MACHINES)
         else:
             raise ValidationError(f"Node name {node_name} is invalid!")
 
@@ -4799,20 +4841,22 @@ def _aks_get_node_name_as(
         cmd,
         node_name,
         managed_resource_group):
-    compute_client = get_compute_client(cmd.cli_ctx)
-
+    from azure.cli.command_modules.vm.aaz.latest.vm import List as VMList
+    from azure.cli.command_modules.vm.operations.vm import VMShow
     if not node_name:
         logger.debug("No node name specified, will randomly select a node from the cluster")
 
-        vm_list = compute_client.virtual_machines.list(managed_resource_group)
+        vm_list = VMList(cli_ctx=cmd.cli_ctx)(command_args={
+            'resource_group': managed_resource_group
+        })
         if not vm_list:
             raise ValidationError(f"No VM found in the managed resource group {managed_resource_group}!")
 
         vm_name = ""
         for vm in vm_list:
             try:
-                _aks_verify_resource(vm, CONST_VIRTUAL_MACHINES)
-                vm_name = vm.name
+                _aks_verify_resource_by_aaz(vm, CONST_VIRTUAL_MACHINES)
+                vm_name = vm.get('name')
                 logger.debug("Select VM: %s", vm_name)
                 break
             except ValidationError as ex:
@@ -4823,10 +4867,14 @@ def _aks_get_node_name_as(
             raise ValidationError("No suitable VM found in the managed resource!")
     else:
         vm_name = node_name
-        vm_info = compute_client.virtual_machines.get(managed_resource_group, vm_name)
+        command_args = {
+            'resource_group': managed_resource_group,
+            'vm_name': vm_name
+        }
+        vm_info = VMShow(cli_ctx=cmd.cli_ctx)(command_args=command_args)
         if not vm_info:
             raise ValidationError(f"VM {vm_name} not found in the managed resource group {managed_resource_group}!")
-        _aks_verify_resource(vm_info, CONST_VIRTUAL_MACHINES)
+        _aks_verify_resource_by_aaz(vm_info, CONST_VIRTUAL_MACHINES)
 
     return vm_name