vm-repair: select a compatible repair VM disk controller

src/vm-repair/HISTORY.rst

@@ -2,6 +2,10 @@
 Release History
 ===============
 
+2.3.1
+++++++
+The repair VM is now created with the SCSI disk controller whenever the selected VM size supports it, even when the source VM uses NVMe. Previously the repair VM inherited the platform default for the size, so a source VM on an NVMe-only size produced an NVMe repair VM. Several repair scripts locate the attached OS disk by its SCSI model name and therefore found no disk on such a repair VM: they completed, reported success, and repaired nothing. Pinning the repair VM to SCSI restores those scripts. When the size only supports NVMe the command now emits a warning instead of failing silently. A new ``--disk-controller-type`` parameter on ``az vm repair create`` overrides the selection.
+
 2.2.6
 ++++++
 Fixing ``az extension add --name vm-repair`` failing with ``Pip failed with status code 2`` on 32-bit Windows installations of the Azure CLI. The extension declared ``opencensus`` as a dependency but never imported it. Because extensions are installed with ``pip --target``, that unused dependency pulled roughly twenty extra packages into the extension folder, including ``cryptography``, which stopped publishing 32-bit Windows wheels in version 49.0.0. On a 32-bit CLI, pip had no compatible wheel, fell back to building ``cryptography`` from source, and failed. Removing the unused ``opencensus`` dependency removes that entire dependency tree.

src/vm-repair/azext_vm_repair/_help.py

@@ -46,6 +46,9 @@
         - name: Create a repair VM with an OS Disk storage type of StandardSSD_LRS.
           text: >
             az vm repair create -g MyResourceGroup -n myVM --repair-username <username> --repair-password <password> --os-disk-type StandardSSD_LRS
+        - name: Create a repair VM using the NVMe disk controller.
+          text: >
+            az vm repair create -g MyResourceGroup -n MySourceVM --disk-controller-type NVMe --verbose
 """
 
 helps['vm repair restore'] = """

src/vm-repair/azext_vm_repair/_params.py

@@ -35,6 +35,7 @@ def load_arguments(self, _):
         c.argument('distro', help='Option to create repair vm from a specific linux distro (rhel7|rhel8|sles12|sles15|ubuntu20|centos7|centos8|oracle7)')
         c.argument('disable_trusted_launch', help='Option to disable Trusted Launch security type on the repair vm by setting the security type to Standard.')
         c.argument('os_disk_type', help='Change the OS Disk storage type from the default of PremiumSSD_LRS to the given value.')
+        c.argument('disk_controller_type', help='Disk controller type for the repair VM (SCSI or NVMe). By default the repair VM uses SCSI when the selected size supports it, so that offline repair scripts can enumerate the attached OS disk.')
         c.argument('tags', help='Quoted string with space-separated key-value pairs in "key=value" format. Will be appended to the tags required for repair resources.')
         c.argument('copy_tags', help='Copy tags from the source VM to the repair VM and its resources.  Can be combined with --tags.')
         c.argument('size', help='The size of the repair VM to create. If not specified, a size matching the source VM will be used.')

src/vm-repair/azext_vm_repair/_validators.py

@@ -33,6 +33,13 @@
 
 
 def validate_create(cmd, namespace):
+    if namespace.disk_controller_type:
+        requested_controller = namespace.disk_controller_type.strip().lower()
+        valid_controllers = {'scsi': 'SCSI', 'nvme': 'NVMe'}
+        if requested_controller not in valid_controllers:
+            raise CLIError("--disk-controller-type must be 'SCSI' or 'NVMe'.")
+        namespace.disk_controller_type = valid_controllers[requested_controller]
+
     check_extension_version(EXTENSION_NAME)
 
     # Check if VM exists and is not classic VM

src/vm-repair/azext_vm_repair/custom.py

@@ -23,6 +23,9 @@
     _call_az_command,
     _clean_up_resources,
     _fetch_compatible_sku,
+    _fetch_source_disk_controller_type,
+    _fetch_sku_disk_controller_types,
+    _select_repair_disk_controller_type,
     _list_resource_ids_in_rg,
     _get_repair_resource_tag,
     _validate_tags_for_command,
@@ -56,7 +59,7 @@
 logger = get_logger(__name__)
 
 
-def create(cmd, vm_name, resource_group_name, repair_password=None, repair_username=None, repair_vm_name=None, copy_disk_name=None, repair_group_name=None, unlock_encrypted_vm=False, enable_nested=False, associate_public_ip=False, distro='ubuntu', encrypt_recovery_key="", disable_trusted_launch=False, os_disk_type=None, tags=None, copy_tags=False, size=None, yes=False):
+def create(cmd, vm_name, resource_group_name, repair_password=None, repair_username=None, repair_vm_name=None, copy_disk_name=None, repair_group_name=None, unlock_encrypted_vm=False, enable_nested=False, associate_public_ip=False, distro='ubuntu', encrypt_recovery_key="", disable_trusted_launch=False, os_disk_type=None, tags=None, copy_tags=False, size=None, disk_controller_type=None, yes=False):
     """
     This function creates a repair VM.
 
@@ -79,6 +82,7 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
     - tags: Tags to apply to the repair VM. Should be a dictionary or a string in key[=value] format.
     - copy_tags: If True, tags will be copied from the source VM to the repair VM. Default is False.
     - size: The size of the repair VM. If not provided, the size of the broken vm will be used.
+    - disk_controller_type: Optional SCSI or NVMe override for the repair VM disk controller.
     """
 
     # Logging all the command parameters, except the sensitive data.
@@ -228,6 +232,16 @@ def create(cmd, vm_name, resource_group_name, repair_password=None, repair_usern
         # Adding the size to the command.
         create_repair_vm_command += ' --size {sku}'.format(sku=sku)
 
+        source_controller = None if disk_controller_type else _fetch_source_disk_controller_type(source_vm)
+        supported_controllers = []
+        if source_controller and str(source_controller).lower() == 'nvme':
+            supported_controllers = _fetch_sku_disk_controller_types(sku, source_vm.location)
+        selected_controller, level, message = _select_repair_disk_controller_type(
+            source_controller, supported_controllers, disk_controller_type)
+        getattr(logger, level)(message)
+        if selected_controller:
+            create_repair_vm_command += ' --disk-controller-type {controller}'.format(controller=selected_controller)
+
         # Setting the availability zone for the repair VM.
         # If the source VM has availability zones, the first one is chosen for the repair VM.
         if source_vm.zones:

src/vm-repair/azext_vm_repair/repair_utils.py

@@ -411,6 +411,43 @@ def _fetch_compatible_sku(source_vm, hyperv, requested_sku=None):
     raise CLIError('Selected VM size: \'{}\' is NOT available in location: \'{}\'.'.format(determined_sku, location))
 
 
+def _fetch_source_disk_controller_type(source_vm):
+    """Return the source VM disk controller type, or None when it is unavailable."""
+    # Older compute SDKs do not model this field, so query ARM through Azure CLI as a fallback.
+    storage_profile = getattr(source_vm, 'storage_profile', None)
+    controller = getattr(storage_profile, 'disk_controller_type', None)
+    if controller:
+        return str(getattr(controller, 'value', controller))
+    vm_id = getattr(source_vm, 'id', None)
+    if not vm_id:
+        return None
+    show_command = 'az vm show --ids {id} --query storageProfile.diskControllerType -o tsv'.format(id=vm_id)
+    return (_call_az_command(show_command) or '').strip() or None
+
+
+def _fetch_sku_disk_controller_types(sku, location):
+    """Return the disk controller types supported by a VM size."""
+    query = "[0].capabilities[?name=='DiskControllerTypes'].value"
+    command = 'az vm list-skus -s {sku} -l {loc} --query "{query}" -o tsv'.format(
+        sku=sku, loc=location, query=query)
+    raw = (_call_az_command(command) or '').strip()
+    return [part.strip() for part in raw.split(',') if part.strip()]
+
+
+def _select_repair_disk_controller_type(source_controller, supported_types, requested=None):
+    """Select a repair VM controller while preserving existing SCSI-based repair scripts."""
+    if requested:
+        return requested, 'info', 'Using requested repair VM disk controller type: {}'.format(requested)
+    if not supported_types:
+        return None, 'debug', 'Could not determine supported disk controller types; using the platform default.'
+    if not source_controller or str(source_controller).lower() != 'nvme':
+        return None, 'debug', 'Source VM is not NVMe; using the platform default for the repair VM size.'
+    normalized_types = {controller.lower(): controller for controller in supported_types}
+    if 'scsi' in normalized_types:
+        return 'SCSI', 'info', 'Source VM uses the NVMe disk controller. Creating the repair VM with SCSI so repair scripts can enumerate the attached OS disk. Override with --disk-controller-type.'
+    return None, 'warning', 'The repair VM size only supports NVMe. Repair scripts that select disks by the SCSI model string will not find the attached OS disk. Use --size to pick a size that supports SCSI, or verify the script handles NVMe.'
+
+
 def _fetch_disk_info(resource_group_name, disk_name):
     """ Returns sku, location, os_type, hyperVgeneration, and disk_id as a tuple """
     show_disk_command = 'az disk show -g {g} -n {name} --query [sku.name,location,osType,hyperVGeneration,id] -o json'.format(g=resource_group_name, name=disk_name)

src/vm-repair/azext_vm_repair/tests/latest/test_disk_controller_type.py

@@ -0,0 +1,110 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+# pylint: disable=protected-access
+import unittest
+from unittest import mock
+
+from knack.util import CLIError
+
+from azext_vm_repair._validators import validate_create
+from azext_vm_repair.repair_utils import (
+    _fetch_sku_disk_controller_types,
+    _fetch_source_disk_controller_type,
+    _select_repair_disk_controller_type,
+)
+
+
+class SelectRepairDiskControllerTypeTest(unittest.TestCase):
+
+    def test_scsi_source_no_flag(self):
+        controller, level, _ = _select_repair_disk_controller_type('SCSI', ['SCSI', 'NVMe'])
+        self.assertIsNone(controller)
+        self.assertEqual('debug', level)
+
+    def test_nvme_source_pins_scsi(self):
+        controller, level, message = _select_repair_disk_controller_type('NVMe', ['SCSI', 'NVMe'])
+        self.assertEqual('SCSI', controller)
+        self.assertEqual('info', level)
+        self.assertIn('enumerate', message)
+
+    def test_nvme_only_size_warns(self):
+        controller, level, message = _select_repair_disk_controller_type('NVMe', ['NVMe'])
+        self.assertIsNone(controller)
+        self.assertEqual('warning', level)
+        self.assertIn('only supports NVMe', message)
+
+    def test_explicit_override_wins(self):
+        controller, level, _ = _select_repair_disk_controller_type('SCSI', ['SCSI', 'NVMe'], 'NVMe')
+        self.assertEqual('NVMe', controller)
+        self.assertEqual('info', level)
+
+    def test_unknown_capability_no_flag(self):
+        controller, level, _ = _select_repair_disk_controller_type('NVMe', [])
+        self.assertIsNone(controller)
+        self.assertEqual('debug', level)
+
+
+class FetchDiskControllerTypeTest(unittest.TestCase):
+
+    def test_source_controller_uses_sdk_value(self):
+        source_vm = mock.MagicMock()
+        source_vm.storage_profile.disk_controller_type = 'NVMe'
+        self.assertEqual('NVMe', _fetch_source_disk_controller_type(source_vm))
+
+    def test_source_controller_uses_sdk_enum_value(self):
+        controller = mock.MagicMock(value='NVMe')
+        controller.__str__.return_value = 'DiskControllerTypes.NVME'
+        source_vm = mock.MagicMock()
+        source_vm.storage_profile.disk_controller_type = controller
+        self.assertEqual('NVMe', _fetch_source_disk_controller_type(source_vm))
+
+    @mock.patch('azext_vm_repair.repair_utils._call_az_command', return_value='SCSI\n')
+    def test_source_controller_falls_back_to_cli(self, mock_call):
+        source_vm = mock.MagicMock()
+        source_vm.storage_profile = mock.MagicMock(spec=[])
+        source_vm.id = '/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/rg/providers/Microsoft.Compute/virtualMachines/vm'
+        self.assertEqual('SCSI', _fetch_source_disk_controller_type(source_vm))
+        self.assertIn('storageProfile.diskControllerType', mock_call.call_args[0][0])
+
+    @mock.patch('azext_vm_repair.repair_utils._call_az_command', return_value='NVMe\n')
+    def test_missing_storage_profile_falls_back_to_cli(self, mock_call):
+        source_vm = mock.MagicMock(spec=['id'])
+        source_vm.id = '/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/rg/providers/Microsoft.Compute/virtualMachines/vm'
+        self.assertEqual('NVMe', _fetch_source_disk_controller_type(source_vm))
+        mock_call.assert_called_once()
+
+    @mock.patch('azext_vm_repair.repair_utils._call_az_command')
+    def test_missing_vm_id_returns_none(self, mock_call):
+        source_vm = mock.MagicMock(spec=[])
+        self.assertIsNone(_fetch_source_disk_controller_type(source_vm))
+        mock_call.assert_not_called()
+
+    @mock.patch('azext_vm_repair.repair_utils._call_az_command', return_value='SCSI,NVMe\n')
+    def test_sku_capabilities_are_parsed(self, _):
+        self.assertEqual(['SCSI', 'NVMe'], _fetch_sku_disk_controller_types('Standard_E2bds_v5', 'westus2'))
+
+
+class ValidateDiskControllerTypeTest(unittest.TestCase):
+
+    def _validate(self, value):
+        namespace = mock.MagicMock()
+        namespace.disk_controller_type = value
+        with mock.patch('azext_vm_repair._validators.check_extension_version', side_effect=RuntimeError):
+            with self.assertRaises(RuntimeError):
+                validate_create(mock.MagicMock(), namespace)
+        return namespace.disk_controller_type
+
+    def test_scsi_is_normalized(self):
+        self.assertEqual('SCSI', self._validate(' scsi '))
+
+    def test_nvme_is_normalized(self):
+        self.assertEqual('NVMe', self._validate('nvme'))
+
+    def test_invalid_value_is_rejected_before_network_call(self):
+        namespace = mock.MagicMock(disk_controller_type='SCSI; rm -rf')
+        with mock.patch('azext_vm_repair._validators.check_extension_version') as mock_version:
+            with self.assertRaises(CLIError):
+                validate_create(mock.MagicMock(), namespace)
+        mock_version.assert_not_called()

src/vm-repair/setup.py

@@ -8,7 +8,7 @@
 from codecs import open
 from setuptools import setup, find_packages
 
-VERSION = "2.2.6"
+VERSION = "2.3.1"
 
 CLASSIFIERS = [
     'Development Status :: 4 - Beta',