From 8f967b1a80be9421eb7171b8ab20784be1368dca Mon Sep 17 00:00:00 2001
From: p3rf Team <p3rf-github-robot@google.com>
Date: Sat, 25 Jul 2026 10:21:26 -0700
Subject: [PATCH] Fix GCE VM creation failure for Gen 3+ machine types.

PiperOrigin-RevId: 953902398
---
 .../providers/gcp/gce_virtual_machine.py      | 16 +++++++++++++-
 tests/gce_virtual_machine_test.py             | 21 +++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/perfkitbenchmarker/providers/gcp/gce_virtual_machine.py b/perfkitbenchmarker/providers/gcp/gce_virtual_machine.py
index ae83b10d80..1b4210c9bc 100644
--- a/perfkitbenchmarker/providers/gcp/gce_virtual_machine.py
+++ b/perfkitbenchmarker/providers/gcp/gce_virtual_machine.py
@@ -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()
@@ -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(
@@ -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]
diff --git a/tests/gce_virtual_machine_test.py b/tests/gce_virtual_machine_test.py
index c27f749783..36fe54dda5 100644
--- a/tests/gce_virtual_machine_test.py
+++ b/tests/gce_virtual_machine_test.py
@@ -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):