Make Azure ARM and SCSI disk support more generic.

PiperOrigin-RevId: 687408232

perfkitbenchmarker/providers/azure/azure_disk.py

@@ -79,17 +79,9 @@
 
 # https://docs.microsoft.com/en-us/azure/virtual-machines/azure-vms-no-temp-disk
 # D/Ev4 and D/E(i)sv4 VMs do not have tmp/OS disk; Dv3, Dsv3, and Ddv4 VMs do.
-# Same for *v5, including Milan and ARM machines.
 AZURE_NO_TMP_DISK_TYPES = [
     r'(Standard_D[0-9]+s?_v4)',
     r'(Standard_E[0-9]+i?s?_v4)',
-    r'(Standard_D[0-9]+l?s?_v5)',
-    r'(Standard_E[0-9]+i?s?_v5)',
-    r'(Standard_D[0-9]+as?_v5)',
-    r'(Standard_E[0-9]+i?as?_v5)',
-    r'(Standard_E[0-9]+i?bs?_v5)',
-    r'(Standard_D[0-9]+pl?s_v5)',
-    r'(Standard_E[0-9]+ps_v5)',
 ]
 
 
@@ -148,6 +140,14 @@ def LocalDriveIsNvme(machine_type):
 
 def HasTempDrive(machine_type):
   """Check if the machine type has the temp drive (sdb)."""
+  # Only applies to some gen 4 VMs.
+  series_number = util.GetMachineSeriesNumber(machine_type)
+  if series_number > 5:
+    # This method is assuming it's a SCSI drive.
+    # TODO(pclay): Support NVMe temp drives in v6 VMs.
+    return False
+  if series_number == 5:
+    return machine_type.endswith('ds_v5')
   return not any(
       re.search(machine_series, machine_type)
       for machine_series in AZURE_NO_TMP_DISK_TYPES

perfkitbenchmarker/providers/azure/azure_virtual_machine.py

@@ -124,24 +124,16 @@
 # Recognized Errors
 _OS_PROVISIONING_TIMED_OUT = 'OSProvisioningTimedOut'
 
-# Map to neoverse-n1
-AZURE_ARM_TYPES = [
-    r'(Standard_D[0-9]+pl?d?s_v5)',
-    r'(Standard_E[0-9]+pd?s_v5)',
-]
-
 CONFIDENTIAL_MILAN_TYPES = [
     r'(Standard_DC[0-9]+as?d?s_v5)',
     r'(Standard_EC[0-9]+as?d?s_v5)',
 ]
 
 # Reference -
 # https://learn.microsoft.com/en-us/azure/virtual-machines/trusted-launch#virtual-machines-sizes
-TRUSTED_LAUNCH_UNSUPPORTED_TYPES = AZURE_ARM_TYPES + [
+TRUSTED_LAUNCH_UNSUPPORTED_TYPES = [
     r'(Standard_A[0-9]+_v2)',
-    r'(Standard_D[0-9]+_v2)',
-    r'(Standard_D[0-9]+_v3)',
-    r'(Standard_E[0-9]+_v3)',
+    r'(Standard_[DE][0-9]+_v[2-9])',
     r'(Standard_M[0-9]+.*)',
     r'(Standard_ND[0-9]+a.*)',
 ]
@@ -617,12 +609,23 @@ def _Exists(self):
       return False
 
 
-def _MachineTypeIsArm(machine_type):
+def _MachineTypeIsArm(machine_type) -> bool:
   """Check if the machine type uses ARM."""
-  return any(
-      re.search(machine_series, machine_type)
-      for machine_series in AZURE_ARM_TYPES
-  )
+  return bool(re.match('Standard_[DE][0-9]+p', machine_type))
+
+
+# TODO(pclay): Move out of this class.
+def _GetArmArch(machine_type) -> str | None:
+  """Get the ARM architecture for the machine type."""
+  if _MachineTypeIsArm(machine_type):
+    generation = util.GetMachineSeriesNumber(machine_type)
+    if generation == 5:
+      return 'neoverse-n1'
+    elif generation == 6:
+      return 'neoverse-n2'
+    else:
+      raise ValueError('Unknown ARM machine type: ' + machine_type)
+  return None
 
 
 class AzureVirtualMachine(
@@ -693,7 +696,7 @@ def __init__(self, vm_spec):
         )
         or self.OS_TYPE in TRUSTED_LAUNCH_UNSUPPORTED_OS_TYPES
     )
-    arm_arch = 'neoverse-n1' if _MachineTypeIsArm(self.machine_type) else None
+    arm_arch = _GetArmArch(self.machine_type)
     if arm_arch:
       self.host_arch = arm_arch
       self.is_aarch64 = True
@@ -1377,14 +1380,20 @@ class Windows2019CoreAzureVirtualMachine(
     BaseWindowsAzureVirtualMachine, windows_virtual_machine.Windows2019CoreMixin
 ):
   GEN2_IMAGE_URN = 'MicrosoftWindowsServer:windowsserver-gen2preview:2019-datacenter-gen2:latest'
-  GEN1_IMAGE_URN = 'MicrosoftWindowsServer:WindowsServer:2019-Datacenter-Core:latest'
+  GEN1_IMAGE_URN = (
+      'MicrosoftWindowsServer:WindowsServer:2019-Datacenter-Core:latest'
+  )
 
 
 class Windows2022CoreAzureVirtualMachine(
     BaseWindowsAzureVirtualMachine, windows_virtual_machine.Windows2022CoreMixin
 ):
-  GEN2_IMAGE_URN = 'MicrosoftWindowsServer:WindowsServer:2022-Datacenter-Core-g2:latest'
-  GEN1_IMAGE_URN = 'MicrosoftWindowsServer:WindowsServer:2022-Datacenter-Core:latest'
+  GEN2_IMAGE_URN = (
+      'MicrosoftWindowsServer:WindowsServer:2022-Datacenter-Core-g2:latest'
+  )
+  GEN1_IMAGE_URN = (
+      'MicrosoftWindowsServer:WindowsServer:2022-Datacenter-Core:latest'
+  )
 
 
 class Windows2016DesktopAzureVirtualMachine(
@@ -1407,7 +1416,9 @@ class Windows2022DesktopAzureVirtualMachine(
     BaseWindowsAzureVirtualMachine,
     windows_virtual_machine.Windows2022DesktopMixin,
 ):
-  GEN2_IMAGE_URN = 'MicrosoftWindowsServer:WindowsServer:2022-Datacenter-g2:latest'
+  GEN2_IMAGE_URN = (
+      'MicrosoftWindowsServer:WindowsServer:2022-Datacenter-g2:latest'
+  )
   GEN1_IMAGE_URN = 'MicrosoftWindowsServer:WindowsServer:2022-Datacenter:latest'
 
 

perfkitbenchmarker/providers/azure/util.py

@@ -264,3 +264,19 @@ def GetMachineFamily(machine_type):
   if match:
     return match.group(1) + match.group(2)
   return None
+
+
+def GetMachineSeriesNumber(machine_type: str) -> int:
+  """Returns the series numberof a machine type.
+
+  Not to be confused with the boot architecture generation number.
+  https://learn.microsoft.com/en-us/azure/virtual-machines/generation-2
+
+  Args:
+    machine_type: Azure machine type
+  """
+  match = re.search(r'_v([0-9]+)$', machine_type)
+  if match:
+    return int(match.group(1))
+  # Azure only adds a v after the first series,
+  return 1

tests/providers/azure/azure_disk_test.py

@@ -17,33 +17,34 @@ def setUp(self):
     with mock.patch.object(azure_disk.AzureDisk, '__init__', lambda self: None):
       self.disk = azure_disk.AzureDisk()
       self.disk.disk_type = 'NOT_LOCAL'
-      self.disk.machine_type = 'fake'
+      self.disk.machine_type = 'fake_v5'
       self.disk.vm = mock.Mock()
       self.disk.vm.SupportsNVMe = mock.Mock()
       self.disk.vm.SupportsNVMe.return_value = False
 
-  def test_get_device_path_starts_at_c(self):
+  def test_device_path_used_to_start_at_c(self):
+    self.disk.machine_type = 'fake_v3'
     self.disk.lun = 0
     self.assertEqual('/dev/sdc', self.disk.GetDevicePath())
 
   def test_get_device_path_eq_z(self):
-    self.disk.lun = 23
+    self.disk.lun = 24
     self.assertEqual('/dev/sdz', self.disk.GetDevicePath())
 
   def test_get_device_path_eq_aa(self):
-    self.disk.lun = 24
+    self.disk.lun = 25
     self.assertEqual('/dev/sdaa', self.disk.GetDevicePath())
 
   def test_get_device_path_eq_ba(self):
-    self.disk.lun = 50
+    self.disk.lun = 51
     self.assertEqual('/dev/sdba', self.disk.GetDevicePath())
 
   def test_get_device_path_greatest_allowable_index(self):
-    self.disk.lun = 699
+    self.disk.lun = 700
     self.assertEqual('/dev/sdzz', self.disk.GetDevicePath())
 
   def test_get_device_path_index_too_large(self):
-    self.disk.lun = 700
+    self.disk.lun = 701
     with self.assertRaises(azure_disk.TooManyAzureDisksError):
       self.disk.GetDevicePath()
 

tests/providers/azure/azure_virtual_machine_test.py

@@ -128,7 +128,7 @@ def setUp(self):
   )
   def testVmCreationError(self, stderror, expected_error):
     spec = azure_virtual_machine.AzureVmSpec(
-        _COMPONENT, machine_type='test_machine_type', zone='testing'
+        _COMPONENT, machine_type='Standard_D2s_v5', zone='testing'
     )
     vm = TestAzureVirtualMachine(spec)
     vm.SetDiskSpec(None, 0)
@@ -139,7 +139,7 @@ def testVmCreationError(self, stderror, expected_error):
   def testInsufficientSpotCapacity(self):
     spec = azure_virtual_machine.AzureVmSpec(
         _COMPONENT,
-        machine_type='test_machine_type',
+        machine_type='Standard_D2s_v5',
         zone='testing',
         low_priority=True,
     )

tests/scratch_disk_test.py

@@ -179,7 +179,7 @@ def _PatchCloudSpecific(self):
 
   def _CreateVm(self):
     vm_spec = azure_virtual_machine.AzureVmSpec(
-        'test_vm_spec.Azure', zone='eastus2', machine_type='test_machine_type'
+        'test_vm_spec.Azure', zone='eastus', machine_type='test_machine_type_v5'
     )
     return azure_virtual_machine.Ubuntu2004BasedAzureVirtualMachine(vm_spec)
 
@@ -201,7 +201,7 @@ def _PatchCloudSpecific(self):
   def _CreateVm(self):
     vm_spec = gce_virtual_machine.GceVmSpec(
         'test_vm_spec.GCP',
-        machine_type='test_machine_type',
+        machine_type='test_machine_type_v5',
         zone='us-central1-a',
     )
     vm = gce_virtual_machine.Ubuntu2004BasedGceVirtualMachine(vm_spec)