Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion perfkitbenchmarker/providers/gcp/gce_virtual_machine.py
Original file line number Diff line number Diff line change
Expand Up @@ -206,7 +206,7 @@ def __init__(self, *args, **kwargs):
self.num_local_ssds: int = None # pyrefly: ignore[bad-assignment]
self.preemptible: bool = None # pyrefly: ignore[bad-assignment]
self.boot_disk_size: int = None # pyrefly: ignore[bad-assignment]
self.boot_disk_type: str = None # pyrefly: ignore[bad-assignment]
self.boot_disk_type: str | None = None
self.boot_disk_iops: int = None # pyrefly: ignore[bad-assignment]
self.boot_disk_throughput: int = None # pyrefly: ignore[bad-assignment]
# But prefer GetProject()
Expand All @@ -223,6 +223,11 @@ def __init__(self, *args, **kwargs):
self.ssd_interface: str
self.mtu: int | None
super().__init__(*args, **kwargs)
# Set default boot disk type for Gen 3 and later machines if not specified.
if self.boot_disk_type is None and isinstance(self.machine_type, str):
generation = _GetMachineGeneration(self.machine_type)
if generation and generation >= 3:
self.boot_disk_type = 'hyperdisk-balanced'
# Copy num_local_ssds from flag values to max_local_disks.
self.max_local_disks: int | None = self.num_local_ssds
self.boot_disk_spec = boot_disk.BootDiskSpec(
Expand Down Expand Up @@ -567,6 +572,15 @@ def GetArmArchitecture(machine_type):
return _MACHINE_TYPE_PREFIX_TO_ARM_ARCH.get(prefix)


def _GetMachineGeneration(machine_type: str) -> int | None:
"""Returns the generation of the machine type (e.g. 3 for c3-standard-4)."""
if not machine_type:
return None
family = machine_type.split('-')[0]
match = re.search(r'\d', family)
return int(match.group()) if match else None


def IsLiveMigratableConfidentialCompute(machine_type: str):
"""Returns True if the provided machine type is both confidential computing and can have its maintenance policy set to MIGRATE."""
family = machine_type.split('-')[0]
Expand Down
21 changes: 21 additions & 0 deletions tests/gce_virtual_machine_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,27 @@ def testCustomMachineTypeFlagOverride(self):
self.assertEqual(result.cpus, 1)
self.assertEqual(result.memory, 7680)

def testDefaultBootDiskTypeGen3Plus(self):
# Gen 3 machine should default to hyperdisk-balanced
result = gce_virtual_machine.GceVmSpec(
_COMPONENT, machine_type='c3-standard-4'
)
self.assertEqual(result.boot_disk_type, 'hyperdisk-balanced')

def testDefaultBootDiskTypeGen2(self):
# Gen 2 machine should NOT default to hyperdisk-balanced (remains None)
result = gce_virtual_machine.GceVmSpec(
_COMPONENT, machine_type='n2-standard-2'
)
self.assertIsNone(result.boot_disk_type)

def testExplicitBootDiskTypeGen3Plus(self):
# Explicit boot_disk_type should be preserved
result = gce_virtual_machine.GceVmSpec(
_COMPONENT, machine_type='c3-standard-4', boot_disk_type='pd-ssd'
)
self.assertEqual(result.boot_disk_type, 'pd-ssd')


class GceVirtualMachineTestCase(pkb_common_test_case.PkbCommonTestCase):

Expand Down