[Data] Make test_dataset_throughput deterministic by increasing workload and applying tolerance

Signed-off-by: dancingactor <[email protected]>

Author: dancingactor

python/ray/data/_internal/stats.py

@@ -1219,6 +1219,16 @@ def to_string(
             out += "\n" + self.runtime_metrics()
 
         return out
+    
+    @property
+    def num_rows_per_s(self) -> float:
+        """Calculates the throughput in rows per second for the entire dataset."""
+        output_num_rows = self.operators_stats[-1].output_num_rows
+        total_num_out_rows = output_num_rows["sum"] if output_num_rows else 0
+        wall_time = self.get_total_wall_time()
+        if not total_num_out_rows or not wall_time:
+            return 0.0
+        return total_num_out_rows / wall_time
 
     @staticmethod
     def _collect_dataset_stats_summaries(
@@ -1368,6 +1378,20 @@ class OperatorStatsSummary:
     node_count: Optional[Dict[str, float]] = None
     task_rows: Optional[Dict[str, float]] = None
 
+    @property
+    def num_rows_per_s(self) -> float:
+        """Calculates the overall throughput in rows per second for this operator."""
+        if not self.output_num_rows or not self.time_total_s:
+            return 0.0
+        return self.output_num_rows["sum"] / self.time_total_s
+
+    @property
+    def num_rows_per_task_s(self) -> float:
+        """Calculates the estimated single-task throughput in rows per second."""
+        if not self.output_num_rows or not self.wall_time:
+            return 0.0
+        return self.output_num_rows["sum"] / self.wall_time["sum"]
+
     @classmethod
     def from_block_metadata(
         cls,

python/ray/data/tests/test_stats.py

@@ -1743,37 +1743,71 @@ def test_task_duration_stats():
         durations, ddof=1
     )  # ddof=1 for sample standard deviation
 
+def test_dataset_throughput_calculation(ray_start_regular_shared):
+    """Test throughput calculations using mock block stats."""
+    from ray.data._internal.stats import DatasetStats
+    from ray.data.block import BlockStats, BlockExecStats
+
+    def create_block_stats(start_time, end_time, num_rows):
+        exec_stats = BlockExecStats()
+        exec_stats.start_time_s = start_time
+        exec_stats.end_time_s = end_time
+        exec_stats.wall_time_s = end_time - start_time
+        exec_stats.cpu_time_s = exec_stats.wall_time_s
+        return BlockStats(num_rows=num_rows, size_bytes=None, exec_stats=exec_stats)
+
+    blocks = [
+        create_block_stats(0.0, 2.0, 100),
+        create_block_stats(0.5, 2.5, 100),
+        create_block_stats(1.0, 3.0, 100),
+    ]
 
-# NOTE: All tests above share a Ray cluster, while the tests below do not. These
-# tests should only be carefully reordered to retain this invariant!
+    stats = DatasetStats(metadata={"Map": blocks}, parent=None)
+    summary = stats.to_summary()
 
+    # Throughput: total rows / total execution duration
+    # Total rows = 300
+    # Duration = max end_time - min start_time = 3.0s
+    # 300 rows / 3s = 100 rows/s
+    assert summary.num_rows_per_s == 100
 
-def test_dataset_throughput(shutdown_only):
-    ray.shutdown()
-    ray.init(num_cpus=2)
 
-    f = dummy_map_batches_sleep(0.01)
-    ds = ray.data.range(100).map(f).materialize().map(f).materialize()
+def test_operator_throughput_calculation(ray_start_regular_shared):
+    """Test operator throughput calculations using mock BlockStats."""
+    from ray.data._internal.stats import OperatorStatsSummary
+    from ray.data.block import BlockStats, BlockExecStats
 
-    operator_pattern = re.compile(
-        r"Operator (\d+).*?\* Operator throughput:\s*.*?\* Ray Data throughput: (\d+\.\d+) rows/s.*?\* Estimated single node throughput: (\d+\.\d+) rows/s",
-        re.DOTALL,
-    )
+    def create_block_stats(start_time, end_time, num_rows):
+        exec_stats = BlockExecStats()
+        exec_stats.start_time_s = start_time
+        exec_stats.end_time_s = end_time
+        exec_stats.wall_time_s = end_time - start_time
+        exec_stats.cpu_time_s = exec_stats.wall_time_s
+
+        return BlockStats(num_rows=num_rows, size_bytes=None, exec_stats=exec_stats)
 
-    # Ray data throughput should always be better than single node throughput for
-    # multi-cpu case.
-    for match in operator_pattern.findall(ds.stats()):
-        assert float(match[1]) >= float(match[2])
+    blocks = [
+        create_block_stats(0.0, 2.0, 100),
+        create_block_stats(0.0, 2.0, 100),
+    ]
 
-    # Pattern to match dataset throughput
-    dataset_pattern = re.compile(
-        r"Dataset throughput:.*?Ray Data throughput: (\d+\.\d+) rows/s.*?Estimated single node throughput: (\d+\.\d+) rows/s",  # noqa: E501
-        re.DOTALL,
+    summary = OperatorStatsSummary.from_block_metadata(
+        operator_name="MockOperator",
+        block_stats=blocks,
+        is_sub_operator=False,
     )
 
-    dataset_match = dataset_pattern.search(ds.stats())
-    assert float(dataset_match[1]) >= float(dataset_match[2])
+    # Total rows = 200
+    # Total operator wall time (from earliest start to latest end) = 2.0s
+    # Sum of individual task wall times = 2.0s + 2.0s = 4.0s
+    # Overall throughput: Total rows / Total operator wall time
+    assert summary.num_rows_per_s == 200 / (2.0 - 0.0)
 
+    # Estimated single task throughput: Total rows / Sum of individual task wall times｀
+    assert summary.num_rows_per_task_s == 200 / (2.0 + 2.0)
+
+# NOTE: All tests above share a Ray cluster, while the tests below do not. These
+# tests should only be carefully reordered to retain this invariant!
 
 def test_individual_operator_num_rows(shutdown_only):
     # The input num rows of an individual operator should be the same as the output num rows of its parent operator.