diff --git a/perfkitbenchmarker/providers/gcp/gce_disk_strategies.py b/perfkitbenchmarker/providers/gcp/gce_disk_strategies.py index 40d2631d80..1b0c593bad 100644 --- a/perfkitbenchmarker/providers/gcp/gce_disk_strategies.py +++ b/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.""" diff --git a/perfkitbenchmarker/providers/gcp/gce_virtual_machine.py b/perfkitbenchmarker/providers/gcp/gce_virtual_machine.py index df2490b019..bfa16aa108 100644 --- a/perfkitbenchmarker/providers/gcp/gce_virtual_machine.py +++ b/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)