Refactored azure virtual machine to use azure disk strategies and ena…

…ble disk provisioning benchmark for Azure

PiperOrigin-RevId: 606310865

CHANGES.next.md

@@ -418,3 +418,5 @@
     create and attach time.
 -   Reduce duplicate code in MaintenanceEventTrigger.AppendSamples().
 -   Attach "run_number" label to "LM Total Time" sample.
+-   Refactored `azure_virtual_machine.py` to use `azure_disk_strategies.py` and
+    Enabled disk provision benchmark for Azure.

perfkitbenchmarker/linux_benchmarks/provisioning_benchmarks/provision_disk_benchmark.py

@@ -27,14 +27,16 @@
 from perfkitbenchmarker import benchmark_spec
 from perfkitbenchmarker import configs
 from perfkitbenchmarker import disk
+from perfkitbenchmarker import errors
 from perfkitbenchmarker import linux_virtual_machine
 from perfkitbenchmarker import sample
+from perfkitbenchmarker.providers.azure import flags as azure_flags
 from perfkitbenchmarker.providers.gcp import flags as gcp_flags
 
+
 FLAGS = flags.FLAGS
 
 BENCHMARK_NAME = 'provision_disk'
-
 BENCHMARK_CONFIG = """
 provision_disk:
   description: >
@@ -45,10 +47,16 @@
         GCP:
           machine_type: n2-standard-2
           zone: us-central1-c
+        Azure:
+          machine_type: Standard_D2s_v5
+          zone: eastus2-2
       disk_spec:
         GCP:
           disk_type: pd-ssd
           disk_size: 10
+        Azure:
+          disk_type: PremiumV2_LRS
+          disk_size: 10
 """
 
 
@@ -70,6 +78,10 @@ def CheckPrerequisites(benchmark_config):
     raise ValueError(
         'gcp_create_disks_with_vm must be set to false for GCP'
     )
+  if FLAGS.cloud == 'Azure' and azure_flags.AZURE_ATTACH_DISK_WITH_CREATE.value:
+    raise errors.Setup.InvalidFlagConfigurationError(
+        'azure_attach_disk_with_create must be set to false for Azure'
+    )
 
 
 def _WaitUntilAttached(vm, dsk) -> None:
@@ -93,7 +105,7 @@ def Run(bm_spec: benchmark_spec.BenchmarkSpec) -> List[sample.Sample]:
         'after VM creation.'
     )
   samples = []
-  disk_details = vm.create_disk_strategy.pd_disk_groups
+  disk_details = vm.create_disk_strategy.remote_disk_groups
   for disk_group in disk_details:
     for disk_details in disk_group:
       total_time = 0

perfkitbenchmarker/providers/azure/azure_disk.py

@@ -38,13 +38,26 @@
 
 MAX_DRIVE_SUFFIX_LENGTH = 2  # Last allowable device is /dev/sdzz.
 
-PREMIUM_STORAGE = 'Premium_LRS'
 PREMIUM_STORAGE_V2 = 'PremiumV2_LRS'
+PREMIUM_STORAGE = 'Premium_LRS'
+PREMIUM_ZRS = 'Premium_ZRS'
+STANDARD_SSD_LRS = 'StandardSSD_LRS'
+STANDARD_SSD_ZRS = 'StandardSSD_ZRS'
 STANDARD_DISK = 'Standard_LRS'
 ULTRA_STORAGE = 'UltraSSD_LRS'
-PREMIUM_ZRS = 'Premium_ZRS'
 
 
+# https://learn.microsoft.com/en-us/rest/api/compute/disks/list?view=rest-compute-2023-10-02&tabs=HTTP#diskstorageaccounttypes
+AZURE_REMOTE_DISK_TYPES = [
+    PREMIUM_STORAGE,
+    PREMIUM_STORAGE_V2,
+    STANDARD_SSD_LRS,
+    STANDARD_SSD_ZRS,
+    STANDARD_DISK,
+    ULTRA_STORAGE,
+    PREMIUM_ZRS,
+]
+
 HOST_CACHING = 'host_caching'
 
 AZURE = 'Azure'

perfkitbenchmarker/providers/azure/azure_disk_strategies.py

@@ -18,20 +18,221 @@
 """
 
 from typing import Any
+
 from absl import flags
+from perfkitbenchmarker import disk
 from perfkitbenchmarker import disk_strategies
+from perfkitbenchmarker import errors
 from perfkitbenchmarker.providers.azure import azure_disk
 
 FLAGS = flags.FLAGS
+virtual_machine = Any  # pylint: disable=invalid-name
 
 
-virtual_machine = Any  # pylint: disable=invalid-name
+def GetCreateDiskStrategy(
+    vm: 'virtual_machine.BaseVirtualMachine',
+    disk_spec: disk.BaseDiskSpec,
+    disk_count: int,
+) -> disk_strategies.CreateDiskStrategy:
+  """Returns the strategies to create disks for the disk type."""
+  if disk_spec and disk_count > 0:
+    if disk_spec.disk_type == disk.LOCAL:
+      return AzureCreateLocalSSDDiskStrategy(vm, disk_spec, disk_count)
+    elif disk_spec.disk_type in azure_disk.AZURE_REMOTE_DISK_TYPES:
+      return AzureCreateRemoteDiskStrategy(vm, disk_spec, disk_count)
+  return AzureCreateNonResourceDiskStrategy(vm, disk_spec, disk_count)
+
+
+class AzureCreateDiskStrategy(disk_strategies.CreateDiskStrategy):
+  """Same as CreateDiskStrategy, but with Base Disk spec."""
+
+  disk_spec: disk.BaseDiskSpec
+
+
+class AzureCreateOSDiskStrategy(AzureCreateDiskStrategy):
+  """Strategies to create OS disk.
+  """
+
+  def __init__(
+      self,
+      vm: 'virtual_machine.BaseVirtualMachine',
+      disk_spec: disk.BaseDiskSpec,
+      disk_count: int,
+  ):
+    super().__init__(vm, disk_spec, disk_count)
+    self.vm = vm
+    self.disk_count = disk_count
+    if disk_spec is None:
+      disk_spec = disk.BaseDiskSpec('azure_os_disk')
+    disk_spec.disk_type = vm.boot_disk_type or vm.storage_account.storage_type
+    if vm.boot_disk_size:
+      disk_spec.disk_size = vm.boot_disk_size
+    self.disk = azure_disk.AzureDisk(
+        self.disk_spec, self.vm, None, is_image=True
+    )
+    self.disk_spec = disk_spec
+
+  def GetCreationCommand(self) -> dict[str, Any]:
+    dic = {}
+    disk_type_cmd = []
+    disk_size_cmd = []
+    if self.disk_spec.disk_type:
+      disk_type_cmd = ['--storage-sku', f'os={self.disk.disk_type}']
+    if self.disk and self.disk.disk_size:
+      disk_size_cmd = ['--os-disk-size-gb', str(self.disk.disk_size)]
+    dic['create-disk'] = disk_type_cmd+disk_size_cmd
+    return dic
+
+
+class AzureCreateRemoteDiskStrategy(AzureCreateDiskStrategy):
+  """Strategies to create remote disks.
+  """
+
+  def __init__(
+      self,
+      vm: 'virtual_machine.BaseVirtualMachine',
+      disk_spec: disk.BaseDiskSpec,
+      disk_count: int,
+  ):
+    super().__init__(vm, disk_spec, disk_count)
+    self.remote_disk_groups = []
+    self.vm = vm
+    for _, disk_spec in enumerate(self.disk_specs):
+      disks = []
+      for _ in range(disk_spec.num_striped_disks):
+        disk_number = (
+            self.vm.remote_disk_counter
+            + 1
+            + self.vm.max_local_disks
+        )
+        self.vm.remote_disk_counter += 1
+        lun = next(self.vm.lun_counter)
+        data_disk = azure_disk.AzureDisk(disk_spec, self.vm, lun)
+        # TODO(user) Clean code to avoid using
+        # disk_number as it is used for windows only
+        data_disk.disk_number = disk_number
+        disks.append(data_disk)
+      self.remote_disk_groups.append(disks)
+
+  def DiskCreatedOnVMCreation(self) -> bool:
+    """Returns whether the disk is created on VM creation."""
+    return False
+
+  def GetSetupDiskStrategy(self) -> disk_strategies.SetUpDiskStrategy:
+    """Returns the SetUpDiskStrategy for the disk."""
+    return AzureSetUpRemoteDiskStrategy(self.vm, self.disk_specs)
+
+
+class AzureCreateLocalSSDDiskStrategy(AzureCreateDiskStrategy):
+  """Strategies to create local disks.
+  """
+
+  def DiskCreatedOnVMCreation(self) -> bool:
+    """Returns whether the disk is created on VM creation."""
+    return True
+
+  def GetSetupDiskStrategy(self) -> disk_strategies.SetUpDiskStrategy:
+    """Returns the SetUpDiskStrategy for the disk."""
+    return AzureSetUpLocalSSDDiskStrategy(self.vm, self.disk_spec)
+
+
+class AzureCreateNonResourceDiskStrategy(
+    disk_strategies.EmptyCreateDiskStrategy
+):
+  """Strategies to create non remote and non local disks like RAM, SMB."""
+
+  def DiskCreatedOnVMCreation(self) -> bool:
+    # This have to be set to False due to _CreateScratchDiskFromDisks in
+    # virtual_machine.py.
+    return False
+
+  def GetSetupDiskStrategy(self) -> disk_strategies.SetUpDiskStrategy:
+    """Returns the SetUpDiskStrategy for the disk."""
+    if not self.disk_spec:
+      return disk_strategies.EmptySetupDiskStrategy(self.vm, self.disk_spec)
+
+    if self.disk_spec.disk_type == disk.RAM:
+      return disk_strategies.SetUpRamDiskStrategy(self.vm, self.disk_spec)
+
+    elif self.disk_spec.disk_type == disk.SMB:
+      return disk_strategies.SetUpSMBDiskStrategy(self.vm, self.disk_spec)
+
+    return disk_strategies.EmptySetupDiskStrategy(self.vm, self.disk_spec)
+
+
+class AzureSetUpLocalSSDDiskStrategy(disk_strategies.SetUpDiskStrategy):
+  """Strategies to setup local disk.
+  """
+
+  def __init__(
+      self,
+      vm: 'virtual_machine.BaseVirtualMachine',
+      disk_spec: disk.BaseDiskSpec,
+  ):
+    super().__init__(vm, disk_spec)
+    self.vm = vm
+    self.disk_spec = disk_spec
+
+  def SetUpDisk(
+      self
+  ):
+    disks = []
+
+    for _ in range(self.disk_spec.num_striped_disks):
+      disk_number = self.vm.local_disk_counter + 1
+      self.vm.local_disk_counter += 1
+      if (
+          self.vm.local_disk_counter
+          > self.vm.max_local_disks
+      ):
+        raise errors.Error('Not enough local disks.')
+      lun = next(self.vm.lun_counter)
+      data_disk = azure_disk.AzureDisk(self.disk_spec, self.vm, lun)
+      data_disk.disk_number = disk_number
+      disks.append(data_disk)
+    if len(disks) > 1:
+      # If the disk_spec called for a striped disk, create one.
+      scratch_disk = disk.StripedDisk(self.disk_spec, disks)
+    else:
+      scratch_disk = disks[0]
+    AzurePrepareScratchDiskStrategy().PrepareScratchDisk(
+        self.vm, scratch_disk, self.disk_spec
+    )
+
+
+class AzureSetUpRemoteDiskStrategy(disk_strategies.SetUpDiskStrategy):
+  """Strategies to setup remote disk.
+  """
+
+  def __init__(
+      self,
+      vm: 'virtual_machine.BaseVirtualMachine',
+      disk_specs: list[disk.BaseDiskSpec],
+  ):
+    super().__init__(vm, disk_specs[0])
+    self.vm = vm
+    self.disk_specs = disk_specs
+
+  def SetUpDisk(self):
+    for disk_spec_id, disk_spec in enumerate(self.disk_specs):
+      disk_group = self.vm.create_disk_strategy.remote_disk_groups[disk_spec_id]
+      if len(disk_group) > 1:
+        # If the disk_spec called for a striped disk, create one.
+        scratch_disk = disk.StripedDisk(disk_spec, disk_group)
+      else:
+        scratch_disk = disk_group[0]
+      if not self.vm.create_disk_strategy.DiskCreatedOnVMCreation():
+        scratch_disk.Create()
+        scratch_disk.Attach(self.vm)
+      AzurePrepareScratchDiskStrategy().PrepareScratchDisk(
+          self.vm, scratch_disk, disk_spec
+      )
 
 
 class AzurePrepareScratchDiskStrategy(
     disk_strategies.PrepareScratchDiskStrategy
 ):
-  """Strategies to prepare scratch disk on AWS."""
+  """Strategies to prepare scratch disk on Azure."""
 
   def GetLocalSSDNames(self) -> list[str]:
     # This is only for Ls-Series local ssd