Create abstraction for disruption event to handle async vs sync operations for Migration

PiperOrigin-RevId: 842572782

Author: raymond13513

perfkitbenchmarker/time_triggers/base_disruption_trigger.py

@@ -64,19 +64,30 @@ class BaseDisruptionTrigger(base_time_trigger.BaseTimeTrigger):
   def __init__(self, delay: int):
     super().__init__(delay)
     self.metadata = {}
-    self.disruption_ends = None
+    self.disruption_events: MutableSequence[DisruptionEvent] = []
 
   def TriggerMethod(self, vm: virtual_machine.VirtualMachine):
-    """Trigger the disruption."""
+    """Trigger the disruption.
+
+    Implementation of this needs to modify the disruption_events list if the
+    operation sync.
+
+    Args:
+      vm: The VirtualMachine to trigger the disruption on.
+    """
     raise NotImplementedError()
 
   def SetUp(self):
     """See base class."""
     raise NotImplementedError()
 
-  def WaitForDisruption(self) -> MutableSequence[DisruptionEvent]:
-    """Wait for disruption to end and return the end time."""
-    return []
+  def WaitForDisruption(self) -> None:
+    """Wait for disruption to end and return the end time.
+
+    Only need to implement this if the operation is async. If the operation is
+    async append the events to the disruption_events list.
+    """
+    pass
 
   def GetMetadataForTrigger(self, event: DisruptionEvent) -> Dict[str, Any]:
     """Get the metadata for the trigger and append it to the samples."""
@@ -97,11 +108,11 @@ def AppendSamples(
     def generate_disruption_total_time_samples() -> (
         MutableSequence[sample.Sample]
     ):
-      events = self.WaitForDisruption()
+      self.WaitForDisruption()
 
       # Host maintenance is in s
       self.disruption_ends = max(
-          [float(d.end_time) * 1000 for d in events],
+          [float(d.end_time) * 1000 for d in self.disruption_events],
           default=0,
       )
 
@@ -120,7 +131,7 @@ def generate_disruption_total_time_samples() -> (
               'seconds',
               sample_metadata | self.GetMetadataForTrigger(d),
           )
-          for d in events
+          for d in self.disruption_events
       ]
 
     samples += generate_disruption_total_time_samples()
@@ -251,7 +262,12 @@ def _AggregateThroughputSample(
 
   def GetDisruptionEnds(self) -> float | None:
     """Get the disruption ends."""
-    return self.disruption_ends
+    if self.disruption_events:
+      return max(
+          [float(d.end_time) * 1000 for d in self.disruption_events],
+          default=None,
+      )
+    return None
 
   def _ComputeLossPercentile(
       self,

perfkitbenchmarker/time_triggers/maintenance_simulation_trigger.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 """Module containning methods for triggering maintenance simulation."""
 
-from collections.abc import MutableSequence
 import datetime
 import json
 import logging
@@ -364,26 +363,13 @@ def SetUp(self):
       for helper in self.gce_simulate_maintenance_helpers.values():
         helper.SetupLMNotification()
 
-  def WaitForDisruption(
-      self,
-  ) -> MutableSequence[base_disruption_trigger.DisruptionEvent]:
+  def WaitForDisruption(self) -> None:
     """Wait for the disruption to end and return the end time."""
     if self.capture_live_migration_timestamps:
       # Block test exit until LM ended.
-      lm_events = []
       for helper in self.gce_simulate_maintenance_helpers.values():
         helper.WaitLMNotificationRelease()
-        lm_events.append(helper.CollectLMNotificationsTime())
-      return lm_events
-    else:
-      return []
-
-  def GetDisruptionEnds(self) -> float | None:
-    """Get the disruption ends."""
-    if self.capture_live_migration_timestamps:
-      # lm ends is computed from LM notification
-      return self.disruption_ends
-    return None
+        self.disruption_events.append(helper.CollectLMNotificationsTime())
 
   @property
   def trigger_name(self) -> str:

tests/time_triggers/base_disruption_trigger_test.py

@@ -49,6 +49,7 @@ class BaseDisruptionTriggerTest(pkb_common_test_case.PkbCommonTestCase):
   def setUp(self):
     super().setUp()
     self.trigger = TestBaseDisruptionTrigger()
+    self.maxDiff = 100000
     self.enter_context(
         mock.patch.object(self.trigger, 'WaitForDisruption', autospec=True)
     )
@@ -63,7 +64,7 @@ def testAppendSamples(self):
     }
     s = []
     vm_spec = mock.MagicMock(spec=benchmark_spec.BenchmarkSpec)
-    self.trigger.WaitForDisruption.return_value = [
+    self.trigger.disruption_events = [
         base_disruption_trigger.DisruptionEvent(
             total_time=10, end_time=10, start_time=4
         )
@@ -79,7 +80,7 @@ def testAppendLossFunctionWithDegradationPercent(self):
     FLAGS.maintenance_degradation_percent = 90
     vm_spec = mock.MagicMock(spec=benchmark_spec.BenchmarkSpec)
     self.trigger.trigger_time = datetime.datetime.fromtimestamp(2)
-    self.trigger.WaitForDisruption.return_value = []
+    self.trigger.disruption_events = []
     self.enter_context(
         mock.patch.object(
             self.trigger, 'GetDisruptionEnds', return_value=None, autospec=True
@@ -208,7 +209,7 @@ def testAppendLossFunctionWithDegradationPercent(self):
   def testAppendLossFunctionWithMissingTimeStampsWithRegression(self):
     vm_spec = mock.MagicMock(spec=benchmark_spec.BenchmarkSpec)
     self.trigger.trigger_time = datetime.datetime.fromtimestamp(2)
-    self.trigger.WaitForDisruption.return_value = []
+    self.trigger.disruption_events = []
     self.enter_context(
         mock.patch.object(
             self.trigger, 'GetDisruptionEnds', return_value=None, autospec=True
@@ -336,7 +337,7 @@ def testAppendLossFunctionWithMissingTimeStampsWithRegression(self):
   @mock.patch.object(time, 'time', mock.MagicMock(return_value=0))
   def testAppendLossFunctionWithMissingTimeStampsNoRegression(self):
     self.trigger.trigger_time = datetime.datetime.fromtimestamp(2)
-    self.trigger.WaitForDisruption.return_value = []
+    self.trigger.disruption_events = []
     self.enter_context(
         mock.patch.object(
             self.trigger, 'GetDisruptionEnds', return_value=None, autospec=True
@@ -465,7 +466,7 @@ def testAppendLossFunctionWithMissingTimeStampsNoRegression(self):
   @mock.patch.object(time, 'time', mock.MagicMock(return_value=0))
   def testAppendLossFunctionSamples(self):
     vm_spec = mock.MagicMock(spec=benchmark_spec.BenchmarkSpec)
-    self.trigger.WaitForDisruption.return_value = []
+    self.trigger.disruption_events = []
     self.enter_context(
         mock.patch.object(
             self.trigger, 'GetDisruptionEnds', return_value=None, autospec=True
@@ -611,7 +612,7 @@ def testAppendLossFunctionSamplesWithDisruptionEnds(self):
     )
     samples = [s]
     self.trigger.trigger_time = datetime.datetime.fromtimestamp(4)
-    self.trigger.WaitForDisruption.return_value = [
+    self.trigger.disruption_events = [
         base_disruption_trigger.DisruptionEvent(
             total_time=100, end_time=8, start_time=4
         )
@@ -811,7 +812,7 @@ def testAppendLossFunctionSamplesContainsMetadata(self):
     )
     samples = [s]
     self.trigger.trigger_time = datetime.datetime.fromtimestamp(4)
-    self.trigger.WaitForDisruption.return_value = [
+    self.trigger.disruption_events = [
         base_disruption_trigger.DisruptionEvent(
             total_time=100, end_time=8, start_time=4
         )
@@ -1016,7 +1017,7 @@ def testAppendLossFunctionSamplesHandleTimeDrift(self):
     )
     samples = [s]
     self.trigger.trigger_time = datetime.datetime.fromtimestamp(4)
-    self.trigger.WaitForDisruption.return_value = [
+    self.trigger.disruption_events = [
         base_disruption_trigger.DisruptionEvent(
             total_time=100, end_time=11, start_time=4
         )
@@ -1175,7 +1176,7 @@ def testMaintenanceEventTriggerAppendSamplesWithMaintenanceDegradationWindow(
     )
     samples = [s]
     self.trigger.trigger_time = datetime.datetime.fromtimestamp(4)
-    self.trigger.WaitForDisruption.return_value = [
+    self.trigger.disruption_events = [
         base_disruption_trigger.DisruptionEvent(
             total_time=100, end_time=11, start_time=4
         )
@@ -1334,7 +1335,7 @@ def testMaintenanceEventTriggerAppendSamplesWithRegressionOutsideMaintenanceWind
     )
     samples = [s]
     self.trigger.trigger_time = datetime.datetime.fromtimestamp(4)
-    self.trigger.WaitForDisruption.return_value = [
+    self.trigger.disruption_events = [
         base_disruption_trigger.DisruptionEvent(
             total_time=100, end_time=8, start_time=4
         )

tests/time_triggers/maintenance_simulation_trigger_test.py

@@ -75,8 +75,9 @@ def testWaitForDisruption(self):
     s = []
     trigger = maintenance_simulation_trigger.MaintenanceEventTrigger()
     self.enter_context(
-        mock.patch.object(trigger, 'WaitForDisruption', return_value=[event])
+        mock.patch.object(trigger, 'WaitForDisruption', return_value=None)
     )
+    trigger.disruption_events = [event]
     trigger.capture_live_migration_timestamps = True
     trigger.vms = [vm]
     trigger.AppendSamples(None, vm_spec, s)
@@ -112,8 +113,9 @@ def testWaitForDisruptionReturnsCorrectValue(self):
     ].CollectLMNotificationsTime.return_value = event
     trigger.capture_live_migration_timestamps = True
     trigger.vms = [vm]
+    trigger.WaitForDisruption()
     self.assertEqual(
-        trigger.WaitForDisruption(),
+        trigger.disruption_events,
         [event],
     )