Move Local Disk creation of GCE to local disk strategy.

PiperOrigin-RevId: 592662695

perfkitbenchmarker/providers/gcp/gce_disk_strategies.py

@@ -20,6 +20,8 @@
 from absl import flags
 from perfkitbenchmarker import disk
 from perfkitbenchmarker import disk_strategies
+from perfkitbenchmarker import errors
+from perfkitbenchmarker import os_types
 from perfkitbenchmarker.providers.gcp import gce_disk
 from perfkitbenchmarker.providers.gcp import util
 
@@ -123,6 +125,71 @@ def DiskCreatedOnVMCreation(self) -> bool:
     return False
 
 
+class SetUpGceLocalDiskStrategy(disk_strategies.SetUpDiskStrategy):
+  """Strategies to set up local disks."""
+
+  def __init__(self, disk_specs: list[disk.BaseDiskSpec]):
+    self.disk_specs = disk_specs
+
+  def SetUpDisk(self, vm: Any, _: disk.BaseDiskSpec):
+    # disk spec is not used here.
+    vm.SetupLocalDisks()
+    for disk_spec in self.disk_specs:
+      disks = []
+      for _ in range(disk_spec.num_striped_disks):
+        if vm.ssd_interface == gce_disk.SCSI:
+          name = 'local-ssd-%d' % vm.local_disk_counter
+          disk_number = vm.local_disk_counter + 1
+        elif vm.ssd_interface == gce_disk.NVME:
+          name = f'local-nvme-ssd-{vm.local_disk_counter}'
+          # TODO(user): Apply for new Gen 3 VMs (barring M3)
+          if vm.machine_type.startswith('c3'):
+            # TODO(user): Also consider removing disk_number
+            disk_number = vm.local_disk_counter
+          else:  # includes N2D CVM
+            disk_number = vm.local_disk_counter + vm.NVME_START_INDEX
+        else:
+          raise errors.Error('Unknown Local SSD Interface.')
+
+        data_disk = gce_disk.GceLocalDisk(disk_spec, name)
+        data_disk.disk_number = disk_number
+        vm.local_disk_counter += 1
+        if vm.local_disk_counter > vm.max_local_disks:
+          raise errors.Error('Not enough local disks.')
+        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
+      )
+      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)
+
+    return scratch_disk
+
+
 class SetUpGcsFuseDiskStrategy(disk_strategies.SetUpDiskStrategy):
   """Strategies to set up ram disks."""
 

perfkitbenchmarker/providers/gcp/gce_virtual_machine.py

@@ -1139,7 +1139,11 @@ def SetupAllScratchDisks(self):
       disk_strategies.SetUpNFSDiskStrategy().SetUpDisk(self, self.disk_specs[0])
       return
     if any((spec.disk_type == disk.LOCAL for spec in self.disk_specs)):
-      self.SetupLocalDisks()
+      gce_disk_strategies.SetUpGceLocalDiskStrategy(self.disk_specs).SetUpDisk(
+          self, self.disk_specs[0]
+      )
+      return
+
     for disk_spec_id, disk_spec in enumerate(self.disk_specs):
       self.CreateScratchDisk(disk_spec_id, disk_spec)
       # TODO(user): Simplify disk logic.
@@ -1160,45 +1164,24 @@ def CreateScratchDisk(self, disk_spec_id, disk_spec):
     disks = []
 
     for i in range(disk_spec.num_striped_disks):
-      if disk_spec.disk_type == disk.LOCAL:
-        name = ''
-        if self.ssd_interface == SCSI:
-          name = 'local-ssd-%d' % self.local_disk_counter
-          disk_number = self.local_disk_counter + 1
-        elif self.ssd_interface == NVME:
-          name = f'local-nvme-ssd-{self.local_disk_counter}'
-          # TODO(user): Apply for new Gen 3 VMs (barring M3)
-          if self.machine_type.startswith('c3'):
-            # TODO(user): Also consider removing disk_number
-            disk_number = self.local_disk_counter
-          else:  # includes N2D CVM
-            disk_number = self.local_disk_counter + self.NVME_START_INDEX
-        else:
-          raise errors.Error('Unknown Local SSD Interface.')
-        data_disk = gce_disk.GceLocalDisk(disk_spec, name)
-        data_disk.disk_number = disk_number
-        self.local_disk_counter += 1
-        if self.local_disk_counter > self.max_local_disks:
-          raise errors.Error('Not enough local disks.')
-      else:  # disk_type is PD
-        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
+      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)