Use class instance variable for setup disk instead of changing the ca…

…ll pattern of SetUpDisk method.

PiperOrigin-RevId: 597929362

perfkitbenchmarker/disk_strategies.py

@@ -110,46 +110,46 @@ def DiskCreatedOnVMCreation(self) -> bool:
 class SetUpDiskStrategy:
   """Strategies to set up ram disks."""
 
-  def SetUpDisk(
+  def __init__(
       self,
-      vm,
+      vm: 'virtual_machine.BaseVirtualMachine',
       disk_spec: disk.BaseDiskSpec,
-  ) -> None:
-    if vm.OS_TYPE in os_types.LINUX_OS_TYPES:
-      self.SetUpDiskOnLinux(vm, disk_spec)
+  ):
+    self.vm = vm
+    self.disk_spec = disk_spec
+
+  def SetUpDisk(self) -> None:
+    if self.vm.OS_TYPE in os_types.LINUX_OS_TYPES:
+      self.SetUpDiskOnLinux()
     else:
-      self.SetUpDiskOnWindows(vm, disk_spec)
+      self.SetUpDiskOnWindows()
 
-  def SetUpDiskOnWindows(self, vm, disk_spec: disk.BaseDiskSpec):
+  def SetUpDiskOnWindows(self):
     """Performs Windows specific setup of ram disk."""
     raise NotImplementedError(
-        f'{disk_spec.disk_type} is not supported on Windows.'
+        f'{self.disk_spec.disk_type} is not supported on Windows.'
     )
 
-  def SetUpDiskOnLinux(self, vm, disk_spec: disk.BaseDiskSpec):
+  def SetUpDiskOnLinux(self):
     """Performs Linux specific setup of ram disk."""
     raise NotImplementedError(
-        f'{disk_spec.disk_type} is not supported on linux.'
+        f'{self.disk_spec.disk_type} is not supported on linux.'
     )
 
 
 class EmptyDiskStrategy(SetUpDiskStrategy):
   """Strategies to set up nothing. This is useful when there is no disk."""
 
-  def SetUpDisk(
-      self,
-      vm,
-      disk_spec: disk.BaseDiskSpec,
-  ) -> None:
-    del vm, disk_spec
+  def SetUpDisk(self) -> None:
+    pass
 
 
 class SetUpRamDiskStrategy(SetUpDiskStrategy):
   """Strategies to set up ram disks."""
 
-  def SetUpDiskOnLinux(self, vm, disk_spec: disk.BaseDiskSpec):
+  def SetUpDiskOnLinux(self):
     """Performs Linux specific setup of ram disk."""
-    scratch_disk = disk.BaseDisk(disk_spec)
+    scratch_disk = disk.BaseDisk(self.disk_spec)
     logging.info(
         'Mounting and creating Ram Disk %s, %s',
         scratch_disk.mount_point,
@@ -159,54 +159,58 @@ def SetUpDiskOnLinux(self, vm, disk_spec: disk.BaseDiskSpec):
         'sudo mkdir -p {0};sudo mount -t tmpfs -o size={1}g tmpfs {0};'
         'sudo chown -R $USER:$USER {0};'
     ).format(scratch_disk.mount_point, scratch_disk.disk_size)
-    vm.RemoteHostCommand(mnt_cmd)
-    vm.scratch_disks.append(disk.BaseDisk(disk_spec))
+    self.vm.RemoteHostCommand(mnt_cmd)
+    self.vm.scratch_disks.append(disk.BaseDisk(self.disk_spec))
 
 
 class SetUpNFSDiskStrategy(SetUpDiskStrategy):
   """Strategies to set up NFS disks."""
 
   def __init__(
-      self, unmanaged_nfs_service: Optional[nfs_service.BaseNfsService] = None
+      self,
+      vm,
+      disk_spec: disk.BaseDiskSpec,
+      unmanaged_nfs_service: Optional[nfs_service.BaseNfsService] = None,
   ):
+    super().__init__(vm, disk_spec)
     if unmanaged_nfs_service:
       self.nfs_service = unmanaged_nfs_service
     else:
       self.nfs_service = getattr(
           pkb_context.GetThreadBenchmarkSpec(), 'nfs_service'
       )
 
-  def SetUpDiskOnLinux(self, vm, disk_spec: disk.BaseDiskSpec):
+  def SetUpDiskOnLinux(self):
     """Performs Linux specific setup of ram disk."""
-    vm.Install('nfs_utils')
+    self.vm.Install('nfs_utils')
     nfs_disk = self.nfs_service.CreateNfsDisk()
-    vm.MountDisk(
+    self.vm.MountDisk(
         nfs_disk.GetDevicePath(),
-        disk_spec.mount_point,
-        disk_spec.disk_type,
+        self.disk_spec.mount_point,
+        self.disk_spec.disk_type,
         nfs_disk.mount_options,
         nfs_disk.fstab_options,
     )
 
-    vm.scratch_disks.append(nfs_disk)
+    self.vm.scratch_disks.append(nfs_disk)
 
 
 class SetUpSMBDiskStrategy(SetUpDiskStrategy):
   """Strategies to set up NFS disks."""
 
-  def SetUpDiskOnLinux(self, vm, disk_spec):
+  def SetUpDiskOnLinux(self):
     """Performs Linux specific setup of ram disk."""
     smb_service = getattr(pkb_context.GetThreadBenchmarkSpec(), 'smb_service')
     smb_disk = smb_service.CreateSmbDisk()
-    vm.MountDisk(
+    self.vm.MountDisk(
         smb_disk.GetDevicePath(),
-        disk_spec.mount_point,
-        disk_spec.disk_type,
+        self.disk_spec.mount_point,
+        self.disk_spec.disk_type,
         smb_disk.mount_options,
         smb_disk.fstab_options,
     )
 
-    vm.scratch_disks.append(smb_disk)
+    self.vm.scratch_disks.append(smb_disk)
 
 
 class PrepareScratchDiskStrategy:

perfkitbenchmarker/providers/gcp/gce_disk_strategies.py

@@ -128,9 +128,6 @@ def DiskCreatedOnVMCreation(self) -> bool:
 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 FindRemoteNVMEDevices(self, nvme_devices):
     """Find the paths for all remote NVME devices inside the VM."""
     remote_nvme_devices = [
@@ -143,7 +140,6 @@ def FindRemoteNVMEDevices(self, nvme_devices):
 
   def UpdateDevicePath(self, scratch_disk, remote_nvme_devices):
     """Updates the paths for all remote NVME devices inside the VM."""
-    disks = []
     if isinstance(scratch_disk, disk.StripedDisk):
       disks = scratch_disk.disks
     else:
@@ -161,55 +157,55 @@ def UpdateDevicePath(self, scratch_disk, remote_nvme_devices):
 class SetUpGceLocalDiskStrategy(SetUpGCEResourceDiskStrategy):
   """Strategies to set up local disks."""
 
-  def SetUpDisk(self, vm: Any, disk_spec: disk.BaseDiskSpec):
+  def SetUpDisk(self):
     # disk spec is not used here.
-    vm.SetupLocalDisks()
+    self.vm.SetupLocalDisks()
     disks = []
-    for _ in range(disk_spec.num_striped_disks):
-      if vm.ssd_interface == gce_disk.SCSI:
-        name = 'local-ssd-%d' % vm.local_disk_counter
-      elif vm.ssd_interface == gce_disk.NVME:
-        name = f'local-nvme-ssd-{vm.local_disk_counter}'
+    for _ in range(self.disk_spec.num_striped_disks):
+      if self.vm.ssd_interface == gce_disk.SCSI:
+        name = 'local-ssd-%d' % self.vm.local_disk_counter
+      elif self.vm.ssd_interface == gce_disk.NVME:
+        name = f'local-nvme-ssd-{self.vm.local_disk_counter}'
       else:
         raise errors.Error('Unknown Local SSD Interface.')
 
-      data_disk = gce_disk.GceLocalDisk(disk_spec, name)
-      vm.local_disk_counter += 1
-      if vm.local_disk_counter > vm.max_local_disks:
+      data_disk = gce_disk.GceLocalDisk(self.disk_spec, name)
+      self.vm.local_disk_counter += 1
+      if self.vm.local_disk_counter > self.vm.max_local_disks:
         raise errors.Error('Not enough local disks.')
       disks.append(data_disk)
 
     if len(disks) == 1:
       scratch_disk = disks[0]
     else:
-      scratch_disk = disk.StripedDisk(disk_spec, disks)
+      scratch_disk = disk.StripedDisk(self.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()
+    if self.vm.OS_TYPE not in os_types.WINDOWS_OS_TYPES:
+      nvme_devices = self.vm.GetNVMEDeviceInfo()
       remote_nvme_devices = self.FindRemoteNVMEDevices(nvme_devices)
       self.UpdateDevicePath(scratch_disk, remote_nvme_devices)
     GCEPrepareScratchDiskStrategy().PrepareScratchDisk(
-        vm, scratch_disk, disk_spec
+        self.vm, scratch_disk, self.disk_spec
     )
 
 
 class SetUpPDDiskStrategy(SetUpGCEResourceDiskStrategy):
   """Strategies to Persistance disk on GCE."""
 
-  def __init__(self, disk_specs: list[gce_disk.GceDiskSpec]):
-    super().__init__(disk_specs)
+  def __init__(self, vm, disk_specs: list[gce_disk.GceDiskSpec]):
+    super().__init__(vm, disk_specs[0])
     self.disk_specs = disk_specs
 
-  def SetUpDisk(self, vm: Any, _: disk.BaseDiskSpec):
+  def SetUpDisk(self):
     # disk spec is not used here.
     for disk_spec_id, disk_spec in enumerate(self.disk_specs):
-      disk_group = vm.create_disk_strategy.pd_disk_groups[disk_spec_id]
+      disk_group = self.vm.create_disk_strategy.pd_disk_groups[disk_spec_id]
       # Create the disk if it is not created on create
-      if not vm.create_disk_strategy.DiskCreatedOnVMCreation():
+      if not self.vm.create_disk_strategy.DiskCreatedOnVMCreation():
         for pd_disk in disk_group:
           pd_disk.Create()
-          pd_disk.Attach(vm)
+          pd_disk.Attach(self.vm)
 
       if len(disk_group) > 1:
         # If the disk_spec called for a striped disk, create one.
@@ -219,12 +215,12 @@ def SetUpDisk(self, vm: Any, _: disk.BaseDiskSpec):
 
       # 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()
+      if self.vm.OS_TYPE not in os_types.WINDOWS_OS_TYPES:
+        nvme_devices = self.vm.GetNVMEDeviceInfo()
         remote_nvme_devices = self.FindRemoteNVMEDevices(nvme_devices)
         self.UpdateDevicePath(scratch_disk, remote_nvme_devices)
       GCEPrepareScratchDiskStrategy().PrepareScratchDisk(
-          vm, scratch_disk, disk_spec
+          self.vm, scratch_disk, disk_spec
       )
 
 
@@ -238,20 +234,20 @@ class SetUpGcsFuseDiskStrategy(disk_strategies.SetUpDiskStrategy):
       'implicit_dirs',
   ]
 
-  def SetUpDiskOnLinux(self, vm, disk_spec):
+  def SetUpDiskOnLinux(self):
     """Performs Linux specific setup of ram disk."""
-    scratch_disk = disk.BaseDisk(disk_spec)
-    vm.Install('gcsfuse')
-    vm.RemoteCommand(
-        f'sudo mkdir -p {disk_spec.mount_point} && sudo chmod a+w'
-        f' {disk_spec.mount_point}'
+    scratch_disk = disk.BaseDisk(self.disk_spec)
+    self.vm.Install('gcsfuse')
+    self.vm.RemoteCommand(
+        f'sudo mkdir -p {self.disk_spec.mount_point} && sudo chmod a+w'
+        f' {self.disk_spec.mount_point}'
     )
 
     opts = ','.join(self.DEFAULT_MOUNT_OPTIONS + FLAGS.mount_options)
     bucket = FLAGS.gcsfuse_bucket
-    target = disk_spec.mount_point
-    vm.RemoteCommand(f'sudo mount -t gcsfuse -o {opts} {bucket} {target}')
-    vm.scratch_disks.append(scratch_disk)
+    target = self.disk_spec.mount_point
+    self.vm.RemoteCommand(f'sudo mount -t gcsfuse -o {opts} {bucket} {target}')
+    self.vm.scratch_disks.append(scratch_disk)
 
 
 class GCEPrepareScratchDiskStrategy(disk_strategies.PrepareScratchDiskStrategy):

perfkitbenchmarker/providers/gcp/gce_virtual_machine.py

@@ -1128,26 +1128,24 @@ def SetupAllScratchDisks(self):
 
     # Prepare vm scratch disks:
     if any((spec.disk_type == disk.RAM for spec in self.disk_specs)):
-      disk_strategies.SetUpRamDiskStrategy().SetUpDisk(self, self.disk_specs[0])
+      disk_strategies.SetUpRamDiskStrategy(self, self.disk_specs[0]).SetUpDisk()
       return
 
     if any((spec.disk_type == disk.OBJECT_STORAGE for spec in self.disk_specs)):
-      gce_disk_strategies.SetUpGcsFuseDiskStrategy().SetUpDisk(
+      gce_disk_strategies.SetUpGcsFuseDiskStrategy(
           self, self.disk_specs[0]
-      )
+      ).SetUpDisk()
       return
     if any((spec.disk_type == disk.NFS for spec in self.disk_specs)):
-      disk_strategies.SetUpNFSDiskStrategy().SetUpDisk(self, self.disk_specs[0])
+      disk_strategies.SetUpNFSDiskStrategy(self, self.disk_specs[0]).SetUpDisk()
       return
     if any((spec.disk_type == disk.LOCAL for spec in self.disk_specs)):
-      gce_disk_strategies.SetUpGceLocalDiskStrategy(self.disk_specs).SetUpDisk(
+      gce_disk_strategies.SetUpGceLocalDiskStrategy(
           self, self.disk_specs[0]
-      )
+      ).SetUpDisk()
       return
 
-    gce_disk_strategies.SetUpPDDiskStrategy(self.disk_specs).SetUpDisk(
-        self, self.disk_specs[0]
-    )
+    gce_disk_strategies.SetUpPDDiskStrategy(self, self.disk_specs).SetUpDisk()
 
   def CreateIpReservation(
       self, ip_address_name: str

perfkitbenchmarker/virtual_machine.py

@@ -1315,14 +1315,14 @@ def SetupAllScratchDisks(self):
     # This method will be depreciate soon.
     # Prepare vm scratch disks:
     if any((spec.disk_type == disk.RAM for spec in self.disk_specs)):
-      disk_strategies.SetUpRamDiskStrategy().SetUpDisk(self, self.disk_specs[0])
+      disk_strategies.SetUpRamDiskStrategy(self, self.disk_specs[0]).SetUpDisk()
       return
     if any((spec.disk_type == disk.NFS for spec in self.disk_specs)):
-      disk_strategies.SetUpNFSDiskStrategy().SetUpDisk(self, self.disk_specs[0])
+      disk_strategies.SetUpNFSDiskStrategy(self, self.disk_specs[0]).SetUpDisk()
       return
 
     if any((spec.disk_type == disk.SMB for spec in self.disk_specs)):
-      disk_strategies.SetUpSMBDiskStrategy().SetUpDisk(self, self.disk_specs[0])
+      disk_strategies.SetUpSMBDiskStrategy(self, self.disk_specs[0]).SetUpDisk()
       return
 
     if any((spec.disk_type == disk.LOCAL for spec in self.disk_specs)):