Add a Step Deviation Query API

PiperOrigin-RevId: 720202588

Author: dipannita08

ml_goodput_measurement/src/goodput.py

@@ -12,9 +12,8 @@
 from cloud_goodput.ml_goodput_measurement.src.checkpoint_badput_calculator import CheckpointBadputCalculator
 from cloud_goodput.ml_goodput_measurement.src.checkpoint_badput_calculator import CheckpointLoggerOptions
 from cloud_goodput.ml_goodput_measurement.src.goodput_cache import GoodputCache
-from cloud_goodput.ml_goodput_measurement.src.goodput_utils import BadputType, GoodputInfo, get_timestamp_from_log_entry
-import numpy as np
-from scipy import stats
+from cloud_goodput.ml_goodput_measurement.src.goodput_utils import BadputType, GoodputInfo, StepInfo
+from cloud_goodput.ml_goodput_measurement.src.goodput_utils import compute_ideal_step_time, get_extra_time_from_anomalous_steps, get_timestamp_from_log_entry
 
 
 _JOB_NAME = 'job_name'
@@ -455,31 +454,6 @@ def _get_current_productive_and_unproductive_time(
       now based on the latest entries retrieved from Cloud Logging.
     """
 
-    def get_extra_time_from_anomalous_steps(step_times: list[Any]) -> float:
-      def get_anomalous_and_normal_step_times(
-          step_times: list[Any],
-      ) -> tuple[list[Any], list[Any]]:
-        mad = stats.median_abs_deviation(step_times)
-        med = np.median(step_times)
-
-        anomalous_step_times = []
-        normal_step_times = []
-        for step_time in step_times:
-          if step_time > (med + mad * 3):
-            anomalous_step_times.append(step_time)
-          else:
-            normal_step_times.append(step_time)
-
-        return anomalous_step_times, normal_step_times
-
-      anomalous_step_times, normal_step_times = (
-          get_anomalous_and_normal_step_times(step_times)
-      )
-      normal_step_mean = np.mean(normal_step_times)
-      return sum(anomalous_step_times) - (
-          len(anomalous_step_times) * normal_step_mean
-      )
-
     def get_segment_productive_and_unproductive_time(
         step_start_data: dict[int, float], curr_step: int
     ) -> tuple[float, dict[BadputType, float]]:
@@ -1010,6 +984,76 @@ def get_job_goodput_interval(
         self._number_of_interruptions,
     )
 
+  def _get_step_times(self):
+    """Helper function to compute step times from the log entries."""
+    step_times = {}
+    previous_step_start_time = None
+    previous_step_count = None
+    for payload in self._current_entries:
+      if _STEP_START_TIME in payload:
+        step_start_time = payload[_STEP_START_TIME]
+        step_count = int(payload[_STEP_COUNT])
+        if (
+            previous_step_start_time is not None
+            and previous_step_count is not None
+            and step_count == previous_step_count + 1
+        ):
+          step_times[previous_step_count] = (
+              step_start_time - previous_step_start_time
+          )
+        previous_step_count = step_count
+        previous_step_start_time = step_start_time
+    return step_times
+
+  def get_step_deviation(
+      self, configured_ideal_step_time: Optional[float] = None
+  ) -> dict[int, float]:
+    """Method to get the step deviation of the current step based on the ideal step time.
+
+    This method computes the ideal step time if one is not provided by the user
+    and returns the step deviation of the current step.
+
+    Args:
+      configured_ideal_step_time: Optional user-defined ideal step time.
+
+    Returns:
+      A dictionary of step deviation for each step.
+    """
+    # Get the log entries.
+    self._update_log_entries()
+    # Compute step times from the log entries.
+    step_times = self._get_step_times()
+    # Get the previous ideal step time from the cache.
+    previous_ideal_step_time = (
+        self._goodput_cache._step_info.ideal_step_time
+        if self._goodput_cache._step_info
+        and self._goodput_cache._step_info.ideal_step_time
+        else None
+    )
+    # Compute ideal step time.
+    ideal_step_time = (
+        configured_ideal_step_time
+        if configured_ideal_step_time is not None
+        else compute_ideal_step_time(
+            step_times=list(step_times.values()),
+            previous_ideal_step_time=previous_ideal_step_time,
+        )
+    )
+
+    # Compute step deviation.
+    step_deviations = {
+        step_count: abs(step_time - ideal_step_time)
+        for step_count, step_time in step_times.items()
+    }
+    # Update the step information in the cache.
+    self._goodput_cache.update_step_info(
+        StepInfo(
+            ideal_step_time=ideal_step_time,
+            step_deviations=step_deviations,
+        )
+    )
+    return step_deviations
+
   def _get_job_badput_breakdown(
       self, total_productive_time, total_unproductive_time, total_job_time
   ):

ml_goodput_measurement/src/goodput_cache.py

@@ -3,7 +3,7 @@
 import datetime
 from typing import Any, Dict, Optional
 
-from cloud_goodput.ml_goodput_measurement.src.goodput_utils import BadputType, GoodputInfo
+from cloud_goodput.ml_goodput_measurement.src.goodput_utils import BadputType, GoodputInfo, StepInfo
 
 _TIME_ENTRY = 'time'
 _JOB_START_TIME = 'job_start_time'
@@ -18,6 +18,11 @@ def __init__(self):
     self._last_entry_timestamp = None
     self._job_start_time = None
     self._job_end_time = None
+    self._step_info = None
+
+  def update_step_info(self, step_info: StepInfo):
+    """Updates the step information."""
+    self._step_info = step_info
 
   def update_cached_entries(self, entries: list[Any]):
     """Updated the cached entries."""
@@ -75,4 +80,3 @@ def clear_cache(self):
   def is_cache_empty(self) -> bool:
     """Checks if the cache is empty."""
     return not self._cached_entries
-

ml_goodput_measurement/src/goodput_utils.py

@@ -3,6 +3,8 @@
 import datetime
 import enum
 from typing import Any, Optional
+import numpy as np
+from scipy import stats
 
 _TIME_ENTRY = 'time'
 
@@ -36,6 +38,62 @@ def __init__(
     self.last_recorded_step = last_recorded_step
 
 
+class StepInfo:
+  """Step Information."""
+
+  def __init__(
+      self,
+      ideal_step_time: float,
+      step_deviations: dict[int, float],
+  ):
+    self.ideal_step_time = ideal_step_time
+    self.step_deviations = step_deviations
+
+
+def compute_ideal_step_time(
+    step_times: list[float], previous_ideal_step_time: Optional[float]
+) -> float:
+  """Helper function to compute the ideal step time."""
+  # Filter out the normal step times from the step times dictionary.
+  mad = stats.median_abs_deviation(step_times)
+  med = np.median(step_times)
+  normal_step_times = []
+  for step_time in step_times:
+    if step_time <= (med + mad * 3):
+      normal_step_times.append(step_time)
+  mean_normal_step_time = np.mean(normal_step_times)
+  if previous_ideal_step_time is not None:
+    return np.mean([mean_normal_step_time, previous_ideal_step_time])
+  return mean_normal_step_time
+
+
+def get_anomalous_and_normal_step_times(
+    step_times: list[Any],
+) -> tuple[list[Any], list[Any]]:
+  mad = stats.median_abs_deviation(step_times)
+  med = np.median(step_times)
+
+  anomalous_step_times = []
+  normal_step_times = []
+  for step_time in step_times:
+    if step_time > (med + mad * 3):
+      anomalous_step_times.append(step_time)
+    else:
+      normal_step_times.append(step_time)
+
+  return anomalous_step_times, normal_step_times
+
+
+def get_extra_time_from_anomalous_steps(step_times: list[Any]) -> float:
+  anomalous_step_times, normal_step_times = get_anomalous_and_normal_step_times(
+      step_times
+  )
+  normal_step_mean = np.mean(normal_step_times)
+  return sum(anomalous_step_times) - (
+      len(anomalous_step_times) * normal_step_mean
+  )
+
+
 def get_timestamp_from_log_entry(
     entry: dict[str, Any],
 ) -> datetime.datetime | None:

ml_goodput_measurement/tests/goodput_test.py

@@ -3,11 +3,13 @@
 import dataclasses
 from dataclasses import asdict
 import datetime
+import random
 import time
 from typing import Optional
 
 from cloud_goodput.ml_goodput_measurement.src import goodput
-from cloud_goodput.ml_goodput_measurement.src.goodput_utils import BadputType, get_timestamp_from_log_entry
+from cloud_goodput.ml_goodput_measurement.src.goodput_utils import BadputType
+from cloud_goodput.ml_goodput_measurement.src.goodput_utils import compute_ideal_step_time, get_timestamp_from_log_entry
 
 from google3.testing.pybase import googletest
 
@@ -1595,6 +1597,52 @@ def test_goodput_badput_with_interval_query(self):
         computed_badput_breakdown[BadputType.WASTED_PROGRESS_FROM_DISRUPTION], 0
     )
 
+  def _generate_step_start_times(self, number_of_steps: int, start_time):
+    """Generate a list of n non-decreasing datetime objects."""
+    max_step_seconds = 600
+    step_start_times = [start_time]
+    for _ in range(1, number_of_steps):
+      increment = random.randint(1, max_step_seconds)
+      new_time = step_start_times[-1] + datetime.timedelta(seconds=increment)
+      step_start_times.append(new_time)
+    return step_start_times
+
+  def test_get_step_deviation(self):
+    """Test function to validate step deviation computation."""
+    job_start_time = datetime.datetime.utcnow()
+    self.goodput_recorder.record_job_start_time(job_start_time)
+    # Generate a list of 100 step start times with random step times.
+    step_count = 0
+    max_steps = 100
+    test_step_start_times = self._generate_step_start_times(
+        number_of_steps=max_steps, start_time=job_start_time
+    )
+
+    # Record step start times.
+    for step_start_time in test_step_start_times:
+      self.goodput_recorder.record_step_start_time(step_count, step_start_time)
+      step_count += 1
+
+    job_end_time = test_step_start_times[-1] + datetime.timedelta(seconds=10)
+    self.goodput_recorder.record_job_end_time(job_end_time)
+
+    step_times = self.goodput_calculator._get_step_times()
+    ideal_step_time = compute_ideal_step_time(
+        step_times=list(step_times.values()), previous_ideal_step_time=None
+    )
+    computed_step_deviations = self.goodput_calculator.get_step_deviation()
+    expected_step_deviations = {
+        step_count: abs(step_time - ideal_step_time)
+        for step_count, step_time in step_times.items()
+    }
+    for step_count, expected_deviation in expected_step_deviations.items():
+      computed_deviation = computed_step_deviations[step_count]
+      self.assertAlmostEqual(
+          expected_deviation,
+          computed_deviation,
+          delta=0.1,
+      )
+
 
 if __name__ == '__main__':
   googletest.main()