Change smart_capacity_retry to retry on a different zone of the same …

…region.

PiperOrigin-RevId: 590743561

perfkitbenchmarker/pkb.py

@@ -1105,20 +1105,27 @@ def _SetOriginalZoneAndFlag(self) -> None:
     if self._GetCurrentZoneFlag() == _ANY_ZONE:
       if pkb_flags.MAX_RETRIES.value < 1:
         FLAGS['retries'].parse(len(self._supported_zones))
-      self._AssignNewZone()
+      self._ChooseAndSetNewZone(self._supported_zones)
     self._original_zone = self._GetCurrentZoneFlag()
     self._original_region = self._utils.GetRegionFromZone(self._original_zone)
 
   def HandleSmartRetries(self, spec: bm_spec.BenchmarkSpec) -> None:
-    """Handles smart zone retry flags if provided."""
+    """Handles smart zone retry flags if provided.
+
+    If quota retry, pick zone in new region. If unsupported or stockout retries,
+    pick zone in same region.
+
+    Args:
+      spec: benchmark spec.
+    """
     if (pkb_flags.SMART_QUOTA_RETRY.value and spec.failed_substatus
         == benchmark_status.FailedSubstatus.QUOTA):
       self._AssignZoneToNewRegion()
     elif (pkb_flags.SMART_CAPACITY_RETRY.value and spec.failed_substatus in {
         benchmark_status.FailedSubstatus.UNSUPPORTED,
         benchmark_status.FailedSubstatus.INSUFFICIENT_CAPACITY
     }):
-      self._AssignNewZone()
+      self._AssignZoneToSameRegion()
 
   def _AssignZoneToNewRegion(self) -> None:
     """Changes zone to be a new zone in the different region."""
@@ -1136,9 +1143,11 @@ def _AssignZoneToNewRegion(self) -> None:
     logging.info('Retry using new region %s', new_region)
     self._ChooseAndSetNewZone(self._utils.GetZonesInRegion(new_region))
 
-  def _AssignNewZone(self) -> None:
-    """Changes zone to be a new zone."""
-    self._ChooseAndSetNewZone(self._supported_zones)
+  def _AssignZoneToSameRegion(self) -> None:
+    """Changes zone to be a new zone in the same region."""
+    supported_zones_in_region = self._utils.GetZonesInRegion(
+        self._original_region).intersection(self._supported_zones)
+    self._ChooseAndSetNewZone(supported_zones_in_region)
 
   def _ChooseAndSetNewZone(self, possible_zones: Set[str]) -> None:
     """Saves the current _zone_flag and sets it to a new zone.

tests/pkb_test.py

@@ -464,7 +464,7 @@ def testRunBenchmarkTask(self, quota_flag_value, capacity_flag_value,
     mock_quota_retry = self.enter_context(
         mock.patch.object(pkb.ZoneRetryManager, '_AssignZoneToNewRegion'))
     mock_capacity_retry = self.enter_context(
-        mock.patch.object(pkb.ZoneRetryManager, '_AssignNewZone'))
+        mock.patch.object(pkb.ZoneRetryManager, '_AssignZoneToSameRegion'))
     self._MockGcpUtils('GetZonesFromMachineType', return_value=set())
 
     benchmark_specs, _ = pkb.RunBenchmarkTask(spec=test_spec)