Move PD SSD on GCE to disk strategy.

PiperOrigin-RevId: 592023608

perfkitbenchmarker/providers/gcp/gce_disk_strategies.py

@@ -125,12 +125,32 @@ def DiskCreatedOnVMCreation(self) -> bool:
     return False
 
 
-class SetUpGceLocalDiskStrategy(disk_strategies.SetUpDiskStrategy):
-  """Strategies to set up local disks."""
+class SetUpGCEResourceDiskStrategy(disk_strategies.SetUpDiskStrategy):
+  """Base Strategy class to set up local ssd and pd-ssd."""
 
   def __init__(self, disk_specs: list[disk.BaseDiskSpec]):
     self.disk_specs = disk_specs
 
+  def _CreateScratchDiskFromDisks(self, vm, disk_spec, disks):
+    """Helper method to create scratch data disks."""
+    # This will soon be moved to disk class. As the disk class should
+    # be able to handle how it's created and attached
+    if len(disks) > 1:
+      # If the disk_spec called for a striped disk, create one.
+      scratch_disk = disk.StripedDisk(disk_spec, disks)
+    else:
+      scratch_disk = disks[0]
+
+    if not vm.create_disk_strategy.DiskCreatedOnVMCreation():
+      scratch_disk.Create()
+      scratch_disk.Attach(vm)
+
+    return scratch_disk
+
+
+class SetUpGceLocalDiskStrategy(SetUpGCEResourceDiskStrategy):
+  """Strategies to set up local disks."""
+
   def SetUpDisk(self, vm: Any, _: disk.BaseDiskSpec):
     # disk spec is not used here.
     vm.SetupLocalDisks()
@@ -173,21 +193,50 @@ def SetUpDisk(self, vm: Any, _: disk.BaseDiskSpec):
       if disk_spec.num_striped_disks > 1:
         break
 
-  def _CreateScratchDiskFromDisks(self, vm, disk_spec, disks):
-    """Helper method to create scratch data disks."""
-    # This will soon be moved to disk class. As the disk class should
-    # be able to handle how it's created and attached
-    if len(disks) > 1:
-      # If the disk_spec called for a striped disk, create one.
-      scratch_disk = disk.StripedDisk(disk_spec, disks)
-    else:
-      scratch_disk = disks[0]
 
-    if not vm.create_disk_strategy.DiskCreatedOnVMCreation():
-      scratch_disk.Create()
-      scratch_disk.Attach(vm)
+class SetUpPDDiskStrategy(SetUpGCEResourceDiskStrategy):
+  """Strategies to Persistance disk on GCE."""
 
-    return scratch_disk
+  def __init__(self, disk_specs: list[gce_disk.GceDiskSpec]):
+    super().__init__(disk_specs)
+    self.disk_specs = disk_specs
+
+  def SetUpDisk(self, vm: Any, _: disk.BaseDiskSpec):
+    # disk spec is not used here.
+    for disk_spec_id, disk_spec in enumerate(self.disk_specs):
+      disks = []
+
+      for i in range(disk_spec.num_striped_disks):
+        name = _GenerateDiskNamePrefix(vm, disk_spec_id, i)
+        data_disk = gce_disk.GceDisk(
+            disk_spec,
+            name,
+            vm.zone,
+            vm.project,
+            replica_zones=disk_spec.replica_zones,
+        )
+        if gce_disk.PdDriveIsNvme(vm):
+          data_disk.interface = gce_disk.NVME
+        else:
+          data_disk.interface = gce_disk.SCSI
+        # Remote disk numbers start at 1+max_local_disks (0 is the system disk
+        # and local disks occupy 1-max_local_disks).
+        data_disk.disk_number = vm.remote_disk_counter + 1 + vm.max_local_disks
+        vm.remote_disk_counter += 1
+        disks.append(data_disk)
+
+      scratch_disk = self._CreateScratchDiskFromDisks(vm, disk_spec, disks)
+      # Device path is needed to stripe disks on Linux, but not on Windows.
+      # The path is not updated for Windows machines.
+      if vm.OS_TYPE not in os_types.WINDOWS_OS_TYPES:
+        nvme_devices = vm.GetNVMEDeviceInfo()
+        remote_nvme_devices = vm.FindRemoteNVMEDevices(
+            scratch_disk, nvme_devices
+        )
+        vm.UpdateDevicePath(scratch_disk, remote_nvme_devices)
+      disk_strategies.PrepareScratchDiskStrategy().PrepareScratchDisk(
+          vm, scratch_disk, disk_spec
+      )
 
 
 class SetUpGcsFuseDiskStrategy(disk_strategies.SetUpDiskStrategy):

perfkitbenchmarker/providers/gcp/gce_virtual_machine.py

@@ -1144,57 +1144,8 @@ def SetupAllScratchDisks(self):
       )
       return
 
-    for disk_spec_id, disk_spec in enumerate(self.disk_specs):
-      self.CreateScratchDisk(disk_spec_id, disk_spec)
-      # TODO(user): Simplify disk logic.
-      if disk_spec.num_striped_disks > 1:
-        # scratch disks has already been created and striped together.
-        break
-    # This must come after Scratch Disk creation to support the
-    # Containerized VM case
-
-  def CreateScratchDisk(self, disk_spec_id, disk_spec):
-    """Create a VM's scratch disk.
-
-    Args:
-      disk_spec_id: Deterministic order of this disk_spec in the VM's list of
-        disk_specs.
-      disk_spec: virtual_machine.BaseDiskSpec object of the disk.
-    """
-    disks = []
-
-    for i in range(disk_spec.num_striped_disks):
-      name = self._GenerateDiskNamePrefix(disk_spec_id, i)
-      data_disk = gce_disk.GceDisk(
-          disk_spec,
-          name,
-          self.zone,
-          self.project,
-          replica_zones=disk_spec.replica_zones,
-      )
-      if gce_disk.PdDriveIsNvme(self):
-        data_disk.interface = gce_disk.NVME
-      else:
-        data_disk.interface = gce_disk.SCSI
-      # Remote disk numbers start at 1+max_local_disks (0 is the system disk
-      # and local disks occupy 1-max_local_disks).
-      data_disk.disk_number = (
-          self.remote_disk_counter + 1 + self.max_local_disks
-      )
-      self.remote_disk_counter += 1
-      disks.append(data_disk)
-
-    scratch_disk = self._CreateScratchDiskFromDisks(disk_spec, disks)
-    # Device path is needed to stripe disks on Linux, but not on Windows.
-    # The path is not updated for Windows machines.
-    if self.OS_TYPE not in os_types.WINDOWS_OS_TYPES:
-      nvme_devices = self.GetNVMEDeviceInfo()
-      remote_nvme_devices = self.FindRemoteNVMEDevices(
-          scratch_disk, nvme_devices
-      )
-      self.UpdateDevicePath(scratch_disk, remote_nvme_devices)
-    disk_strategies.PrepareScratchDiskStrategy().PrepareScratchDisk(
-        self, scratch_disk, disk_spec
+    gce_disk_strategies.SetUpPDDiskStrategy(self.disk_specs).SetUpDisk(
+        self, self.disk_specs[0]
     )
 
   def CreateIpReservation(
@@ -1254,22 +1205,6 @@ def UpdateDevicePath(self, scratch_disk, remote_nvme_devices):
       if d.disk_type in gce_disk.GCE_REMOTE_DISK_TYPES and d.interface == NVME:
         d.name = remote_nvme_devices.pop()
 
-  def _CreateScratchDiskFromDisks(self, disk_spec, disks):
-    """Helper method to create scratch data disks."""
-    # This will soon be moved to disk class. As the disk class should
-    # be able to handle how it's created and attached
-    if len(disks) > 1:
-      # If the disk_spec called for a striped disk, create one.
-      scratch_disk = disk.StripedDisk(disk_spec, disks)
-    else:
-      scratch_disk = disks[0]
-
-    if not self.create_disk_strategy.DiskCreatedOnVMCreation():
-      scratch_disk.Create()
-      scratch_disk.Attach(self)
-
-    return scratch_disk
-
   def AddMetadata(self, **kwargs):
     """Adds metadata to disk."""
     # vm metadata added to vm on creation.

tests/scratch_disk_test.py

@@ -109,24 +109,40 @@ def testScratchDisks(self):
     """
 
     vm = self._CreateVm()
-    disk_spec = self.GetDiskSpec(mount_point='/mountpoint0')
-    vm.SetDiskSpec(disk_spec, 1)
-    vm.CreateScratchDisk(0, disk_spec)
+    disk_spec = self.GetDiskSpec(mount_point='/mountpoint')
+    vm.SetDiskSpec(disk_spec, 2)
 
-    assert len(vm.scratch_disks) == 1, 'Disk not added to scratch disks.'
-
-    scratch_disk = vm.scratch_disks[0]
-
-    scratch_disk.Create.assert_called_once_with()
-    vm.FormatDisk.assert_called_once_with(scratch_disk.GetDevicePath(), None)
-    vm.MountDisk.assert_called_once_with(
-        scratch_disk.GetDevicePath(), '/mountpoint0',
-        None, scratch_disk.mount_options, scratch_disk.fstab_options)
+    vm.SetupAllScratchDisks()
+    assert len(vm.scratch_disks) == 2, 'Disk not added to scratch disks.'
 
-    disk_spec = self.GetDiskSpec(mount_point='/mountpoint1')
-    vm.CreateScratchDisk(0, disk_spec)
+    scratch_disk1 = vm.scratch_disks[0]
+    scratch_disk2 = vm.scratch_disks[1]
+    scratch_disk1.Create.assert_called_once_with()
+    scratch_disk2.Create.assert_called_once_with()
+    format_disk_callls = [
+        mock.call(scratch_disk1.GetDevicePath(), None),
+        mock.call(scratch_disk2.GetDevicePath(), None),
+    ]
 
-    assert len(vm.scratch_disks) == 2, 'Disk not added to scratch disks.'
+    vm.FormatDisk.assert_has_calls(format_disk_callls, None)
+
+    mount_disk_callls = [
+        mock.call(
+            scratch_disk1.GetDevicePath(),
+            '/mountpoint0',
+            None,
+            scratch_disk1.mount_options,
+            scratch_disk1.fstab_options,
+        ),
+        mock.call(
+            scratch_disk2.GetDevicePath(),
+            '/mountpoint1',
+            None,
+            scratch_disk2.mount_options,
+            scratch_disk2.fstab_options,
+        ),
+    ]
+    vm.MountDisk.assert_has_calls(mount_disk_callls, None)
 
     # Check that these execute without exceptions. The return value
     # is a MagicMock, not a string, so we can't compare to expected results.
@@ -136,14 +152,6 @@ def testScratchDisks(self):
     with self.assertRaises(errors.Error):
       vm.GetScratchDir(2)
 
-    scratch_disk = vm.scratch_disks[1]
-
-    scratch_disk.Create.assert_called_once_with()
-    vm.FormatDisk.assert_called_with(scratch_disk.GetDevicePath(), None)
-    vm.MountDisk.assert_called_with(
-        scratch_disk.GetDevicePath(), '/mountpoint1',
-        None, scratch_disk.mount_options, scratch_disk.fstab_options)
-
     vm.DeleteScratchDisks()
 
     vm.scratch_disks[0].Delete.assert_called_once_with()