chore: import upstream snapshot with attribution

python/ray/data/tests/test_op_runtime_metrics.py

@@ -0,0 +1,494 @@
+import time
+from unittest.mock import MagicMock, patch
+
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data._internal.execution.interfaces import BlockEntry, RefBundle
+from ray.data._internal.execution.interfaces.op_runtime_metrics import (
+    OpRuntimeMetrics,
+)
+from ray.data._internal.execution.interfaces.physical_operator import (
+    TaskExecDriverStats,
+)
+from ray.data._internal.util import KiB
+from ray.data.block import BlockExecStats, BlockMetadata, TaskExecWorkerStats
+from ray.data.context import DataContext
+
+
+def test_average_max_uss_per_task():
+    op = MagicMock()
+    op.data_context.enable_get_object_locations_for_metrics = False
+    metrics = OpRuntimeMetrics(op)
+    assert metrics.average_max_uss_per_task is None
+
+    input_bundle = RefBundle([], owns_blocks=False, schema=None)
+
+    # Submit and finish first task with USS of 100 bytes.
+    metrics.on_task_submitted(0, input_bundle)
+    metrics.on_task_finished(
+        0,
+        None,
+        TaskExecWorkerStats(task_wall_time_s=1.0, max_uss_bytes=100),
+        TaskExecDriverStats(task_output_backpressure_s=0),
+    )
+    assert metrics.average_max_uss_per_task == 100
+
+    # Submit and finish second task with USS of 300 bytes.
+    metrics.on_task_submitted(1, input_bundle)
+    metrics.on_task_finished(
+        1,
+        None,
+        TaskExecWorkerStats(task_wall_time_s=1.0, max_uss_bytes=300),
+        TaskExecDriverStats(task_output_backpressure_s=0),
+    )
+    assert metrics.average_max_uss_per_task == 200  # (100 + 300) / 2
+
+
+def test_task_completion_time_histogram():
+    """Test task completion time histogram bucket assignment and counting."""
+    op = MagicMock()
+    op.data_context.enable_get_object_locations_for_metrics = False
+    metrics = OpRuntimeMetrics(op)
+
+    # Test different completion times
+    # Buckets: [0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 7.5, 10.0, 15.0, 20.0, 25.0, 50.0, 75.0, 100.0, 150.0, 500.0, 1000.0, 2500.0, 5000.0]
+    test_cases = [
+        (0.05, 0),  # Very fast task (0.05s) - should go to first bucket (0.1)
+        (0.2, 1),  # Fast task (0.2s) - should go to second bucket (0.25)
+        (0.6, 3),  # Medium task (0.6s) - should go to fourth bucket (1.0)
+        (1.5, 4),  # Slower task (1.5s) - should go to fifth bucket (2.5)
+        (3.0, 5),  # Slow task (3.0s) - should go to sixth bucket (5.0)
+    ]
+
+    for i, (completion_time, expected_bucket) in enumerate(test_cases):
+        # Create input bundle
+        input_bundle = RefBundle([], owns_blocks=False, schema=None)
+
+        # Submit task (this will create the RunningTaskInfo with current time)
+        metrics.on_task_submitted(i, input_bundle)
+
+        # Manually adjust the start time to simulate the completion time
+        metrics._running_tasks[i].start_time = time.perf_counter() - completion_time
+
+        # Complete the task
+        metrics.on_task_finished(
+            i,
+            None,
+            TaskExecWorkerStats(task_wall_time_s=completion_time),
+            TaskExecDriverStats(task_output_backpressure_s=0),
+        )
+
+        # Check that the correct bucket was incremented
+        assert metrics.task_completion_time._bucket_counts[expected_bucket] == 1
+
+        # Reset for next test
+        metrics.task_completion_time._bucket_counts[expected_bucket] = 0
+
+
+def test_block_completion_time_histogram():
+    """Test block completion time histogram bucket assignment and counting.
+
+    Block completion time = (cum_block_gen_time_s + cum_block_ser_time_s) / num_outputs
+    """
+    op = MagicMock()
+    op.data_context.enable_get_object_locations_for_metrics = False
+    metrics = OpRuntimeMetrics(op)
+
+    # Test different block generation scenarios
+    # Buckets: [0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 7.5, 10.0, 15.0, 20.0, 25.0, 50.0, 75.0, 100.0, 150.0, 500.0, 1000.0, 2500.0, 5000.0]
+    # Each test case: (num_blocks, gen_time, ser_time, expected_bucket)
+    # Per-block time = (gen_time + ser_time) / num_blocks
+    test_cases = [
+        # 1 block, 0.08s gen + 0.02s ser = 0.1s total -> 0.1s per block -> bucket 0 (0.1)
+        (1, 0.08, 0.02, 0),
+        # 2 blocks, 0.4s gen + 0.1s ser = 0.5s total -> 0.25s per block -> bucket 1 (0.25)
+        (2, 0.4, 0.1, 1),
+        # 1 block, 0.5s gen + 0.1s ser = 0.6s total -> 0.6s per block -> bucket 3 (1.0)
+        (1, 0.5, 0.1, 3),
+        # 3 blocks, 1.2s gen + 0.3s ser = 1.5s total -> 0.5s per block -> bucket 2 (0.5)
+        (3, 1.2, 0.3, 2),
+    ]
+
+    for i, (num_blocks, gen_time, ser_time, expected_bucket) in enumerate(test_cases):
+        # Create input bundle
+        input_bundle = RefBundle([], owns_blocks=False, schema=None)
+
+        # Submit task
+        metrics.on_task_submitted(i, input_bundle)
+
+        # Manually set the task info to simulate the block generation
+        metrics._running_tasks[i].num_outputs = num_blocks
+        metrics._running_tasks[i].cum_block_gen_time_s = gen_time
+        metrics._running_tasks[i].cum_block_ser_time_s = ser_time
+
+        # Complete the task
+        metrics.on_task_finished(
+            i,
+            None,
+            TaskExecWorkerStats(task_wall_time_s=gen_time + ser_time),
+            TaskExecDriverStats(task_output_backpressure_s=0),
+        )
+
+        # Check that the correct bucket was incremented by the number of blocks
+        assert (
+            metrics.block_completion_time._bucket_counts[expected_bucket] == num_blocks
+        )
+
+        # Reset for next test
+        metrics.block_completion_time._bucket_counts[expected_bucket] = 0
+
+
+@patch("time.perf_counter")
+def test_task_completion_time_excl_backpressure(mock_perf_counter):
+    """Test that average_task_completion_time_excl_backpressure_s correctly
+    subtracts output backpressure from the driver's wall-clock task time.
+
+    Scheduling time is estimated as the time from task submission to the first
+    output arriving on the driver, minus the worker-side time to generate and
+    serialize that first block.
+    """
+    op = MagicMock()
+    op.data_context.enable_get_object_locations_for_metrics = False
+
+    metrics = OpRuntimeMetrics(op)
+
+    test_cases = [
+        # (driver_wall_time_s, scheduling_time_s, backpressure_time_s, gen_time_s, ser_time_s, num_outputs)
+        (2.0, 0.2, 0.5, 0.25, 0.05, 2),  # Task 0
+        (1.5, 0.2, 0.2, 0.3, 0.05, 1),  # Task 1
+        (3.0, 0.2, 1.0, 0.3, 0.05, 3),  # Task 2
+    ]
+
+    def create_output_bundle(gen_time_s, ser_time_s):
+        block = ray.put(pa.Table.from_pydict({}))
+        stats = BlockExecStats(
+            wall_time_s=gen_time_s,
+            block_ser_time_s=ser_time_s,
+        )
+        metadata = BlockMetadata(
+            num_rows=1,
+            size_bytes=0,
+            input_files=None,
+            exec_stats=stats,
+        )
+        return RefBundle([BlockEntry(block, metadata)], owns_blocks=False, schema=None)
+
+    total_gen_ser = 0
+    cumulative_scheduling_time_s = 0.0
+    clock = 0.0
+    for i, tc in enumerate(test_cases):
+        (
+            driver_wall_time_s,
+            scheduling_time_s,
+            output_bp_time_s,
+            gen_time_s,
+            ser_time_s,
+            num_outputs,
+        ) = tc
+
+        input_bundle = RefBundle([], owns_blocks=False, schema=None)
+
+        # Freeze time at task submission
+        submit_time = clock
+        mock_perf_counter.return_value = clock
+        metrics.on_task_submitted(i, input_bundle)
+
+        # Advance clock to first output arrival on driver:
+        #   time_to_first_block = scheduling + gen + ser
+        clock = submit_time + scheduling_time_s + gen_time_s + ser_time_s
+        mock_perf_counter.return_value = clock
+        metrics.on_task_output_generated(
+            i, create_output_bundle(gen_time_s, ser_time_s)
+        )
+
+        # Verify that average_task_scheduling_time_s is correct *before* the
+        # task finishes.  The numerator (task_scheduling_time_s) is incremented
+        # on first output, so the denominator must be num_tasks_have_outputs
+        # (not num_tasks_finished) for the average to be accurate mid-flight.
+        cumulative_scheduling_time_s += scheduling_time_s
+        num_tasks_with_output = i + 1
+        assert metrics.average_task_scheduling_time_s == pytest.approx(
+            cumulative_scheduling_time_s / num_tasks_with_output
+        )
+
+        # Generate remaining outputs (won't affect scheduling time)
+        for _ in range(num_outputs - 1):
+            clock += gen_time_s + ser_time_s
+            mock_perf_counter.return_value = clock
+            metrics.on_task_output_generated(
+                i, create_output_bundle(gen_time_s, ser_time_s)
+            )
+
+        total_gen_ser += num_outputs * (gen_time_s + ser_time_s)
+
+        # Advance clock to task finish
+        clock = submit_time + driver_wall_time_s
+        mock_perf_counter.return_value = clock
+
+        metrics.on_task_finished(
+            i,
+            None,
+            TaskExecWorkerStats(
+                task_wall_time_s=driver_wall_time_s - scheduling_time_s
+            ),
+            TaskExecDriverStats(task_output_backpressure_s=output_bp_time_s),
+        )
+
+    num_tasks = len(test_cases)
+
+    total_driver_wall_time_s = sum(t[0] for t in test_cases)
+    total_scheduling_time_s = sum(t[1] for t in test_cases)
+    total_output_bp_time_s = sum(t[2] for t in test_cases)
+
+    total_worker_wall_time_s = sum(
+        t[0] - t[1] for t in test_cases  # driver_wall_time_s - scheduling_time_s
+    )
+
+    # Raw counters
+    assert metrics.task_block_gen_and_ser_time_s == pytest.approx(total_gen_ser)
+    assert metrics.task_completion_time_s == pytest.approx(total_driver_wall_time_s)
+    assert metrics.task_worker_completion_time_s == pytest.approx(
+        total_worker_wall_time_s
+    )
+    assert metrics.task_scheduling_time_s == pytest.approx(total_scheduling_time_s)
+    assert metrics.task_output_backpressure_time_s == pytest.approx(
+        total_output_bp_time_s
+    )
+
+    # Derived averages
+    assert metrics.average_total_task_completion_time_s == pytest.approx(
+        total_driver_wall_time_s / num_tasks
+    )
+    assert metrics.average_task_scheduling_time_s == pytest.approx(
+        total_scheduling_time_s / num_tasks  # all tasks produced output
+    )
+    assert metrics.average_task_output_backpressure_time_s == pytest.approx(
+        total_output_bp_time_s / num_tasks
+    )
+    assert metrics.average_task_completion_time_excl_backpressure_s == pytest.approx(
+        (total_driver_wall_time_s - total_output_bp_time_s) / num_tasks
+    )
+
+
+def test_block_size_bytes_histogram():
+    """Test block size bytes histogram bucket assignment and counting."""
+    op = MagicMock()
+    op.data_context.enable_get_object_locations_for_metrics = False
+    metrics = OpRuntimeMetrics(op)
+
+    def create_bundle_with_size(size_bytes):
+        block = ray.put(pa.Table.from_pydict({}))
+        stats = BlockExecStats(
+            wall_time_s=0,
+            block_ser_time_s=0,
+        )
+        metadata = BlockMetadata(
+            num_rows=0,
+            size_bytes=size_bytes,
+            input_files=None,
+            exec_stats=stats,
+        )
+        return RefBundle([BlockEntry(block, metadata)], owns_blocks=False, schema=None)
+
+    # Test different block sizes
+    # Buckets: [1KB, 8KB, 64KB, 128KB, 256KB, 512KB, 1MB, 8MB, 64MB, 128MB, 256MB, 512MB, 1GB, 4GB, 16GB, 64GB, 128GB, 256GB, 512GB, 1024GB, 4096GB]
+    test_cases = [
+        (512, 0),  # 512 bytes -> first bucket (1KB)
+        (2 * KiB, 1),  # 2 KiB -> second bucket (8KB)
+        (32 * KiB, 2),  # 32 KiB -> third bucket (64KB)
+        (100 * KiB, 3),  # 100 KiB -> fourth bucket (128KB)
+        (500 * KiB, 5),  # 500 KiB -> sixth bucket (512KB)
+    ]
+
+    for i, (size_bytes, expected_bucket) in enumerate(test_cases):
+        # Create input bundle (can be empty for this test)
+        input_bundle = RefBundle([], owns_blocks=False, schema=None)
+
+        # Submit task
+        metrics.on_task_submitted(i, input_bundle)
+
+        # Create output bundle with the size we want to test
+        output_bundle = create_bundle_with_size(size_bytes)
+
+        # Generate output
+        metrics.on_task_output_generated(i, output_bundle)
+
+        # Check that the correct bucket was incremented
+        assert metrics.block_size_bytes._bucket_counts[expected_bucket] == 1
+
+        # Reset for next test
+        metrics.block_size_bytes._bucket_counts[expected_bucket] = 0
+
+
+def test_block_size_rows_histogram():
+    """Test block size rows histogram bucket assignment and counting."""
+    op = MagicMock()
+    op.data_context.enable_get_object_locations_for_metrics = False
+    metrics = OpRuntimeMetrics(op)
+
+    def create_bundle_with_rows(num_rows):
+        block = ray.put(pa.Table.from_pydict({}))
+        stats = BlockExecStats(
+            wall_time_s=0,
+            block_ser_time_s=0,
+        )
+        metadata = BlockMetadata(
+            num_rows=num_rows,
+            size_bytes=0,
+            input_files=None,
+            exec_stats=stats,
+        )
+        return RefBundle([BlockEntry(block, metadata)], owns_blocks=False, schema=None)
+
+    # Test different row counts
+    # Buckets: [1, 5, 10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10000, 25000, 50000, 100000, 250000, 500000, 1000000, 2500000, 5000000, 10000000]
+    test_cases = [
+        (1, 0),  # 1 row -> first bucket (1)
+        (3, 1),  # 3 rows -> second bucket (5)
+        (7, 2),  # 7 rows -> third bucket (10)
+        (15, 3),  # 15 rows -> fourth bucket (25)
+        (30, 4),  # 30 rows -> fifth bucket (50)
+        (75, 5),  # 75 rows -> sixth bucket (100)
+    ]
+
+    for i, (num_rows, expected_bucket) in enumerate(test_cases):
+        # Create input bundle (can be empty for this test)
+        input_bundle = RefBundle([], owns_blocks=False, schema=None)
+
+        # Submit task
+        metrics.on_task_submitted(i, input_bundle)
+
+        # Create output bundle with the row count we want to test
+        output_bundle = create_bundle_with_rows(num_rows)
+
+        # Generate output
+        metrics.on_task_output_generated(i, output_bundle)
+
+        # Check that the correct bucket was incremented
+        assert metrics.block_size_rows._bucket_counts[expected_bucket] == 1
+
+        # Reset for next test
+        metrics.block_size_rows._bucket_counts[expected_bucket] = 0
+
+
+@pytest.fixture
+def metrics_config_no_sample_with_target(restore_data_context):  # noqa: F811
+    """Fixture for no-sample scenario with target_max_block_size set."""
+    ctx = DataContext.get_current()
+    ctx.target_max_block_size = 128 * 1024 * 1024  # 128MB
+    ctx._max_num_blocks_in_streaming_gen_buffer = 2
+
+    op = MagicMock()
+    op.data_context = ctx
+    metrics = OpRuntimeMetrics(op)
+    return metrics
+
+
+@pytest.fixture
+def metrics_config_no_sample_with_none(restore_data_context):  # noqa: F811
+    """Fixture for no-sample scenario with target_max_block_size=None."""
+    ctx = DataContext.get_current()
+    ctx.target_max_block_size = None
+    ctx._max_num_blocks_in_streaming_gen_buffer = 1
+
+    op = MagicMock()
+    op.data_context = ctx
+    metrics = OpRuntimeMetrics(op)
+    return metrics
+
+
+@pytest.fixture
+def metrics_config_with_sample(restore_data_context):  # noqa: F811
+    """Fixture for scenario with average_bytes_per_output available."""
+    ctx = DataContext.get_current()
+    ctx.target_max_block_size = 128 * 1024 * 1024  # 128MB
+    ctx._max_num_blocks_in_streaming_gen_buffer = 1
+
+    op = MagicMock()
+    op.data_context = ctx
+    metrics = OpRuntimeMetrics(op)
+
+    # Simulate having samples: set bytes_task_outputs_generated and
+    # num_task_outputs_generated to make average_bytes_per_output available
+    actual_block_size = 150 * 1024 * 1024  # 150MB
+    metrics.bytes_task_outputs_generated = actual_block_size
+    metrics.num_task_outputs_generated = 1
+
+    return metrics
+
+
+@pytest.fixture
+def metrics_config_pending_outputs_no_sample(
+    restore_data_context,  # noqa: F811
+):
+    """Fixture for pending outputs during no-sample with target set."""
+    ctx = DataContext.get_current()
+    ctx.target_max_block_size = 64 * 1024 * 1024  # 64MB
+    ctx._max_num_blocks_in_streaming_gen_buffer = 2
+
+    op = MagicMock()
+    op.data_context = ctx
+    metrics = OpRuntimeMetrics(op)
+    metrics.num_tasks_running = 3
+    return metrics
+
+
+@pytest.fixture
+def metrics_config_pending_outputs_none(restore_data_context):  # noqa: F811
+    """Fixture for pending outputs during no-sample with target=None."""
+    ctx = DataContext.get_current()
+    ctx.target_max_block_size = None
+    ctx._max_num_blocks_in_streaming_gen_buffer = 1
+
+    op = MagicMock()
+    op.data_context = ctx
+    metrics = OpRuntimeMetrics(op)
+    metrics.num_tasks_running = 2
+    return metrics
+
+
+@pytest.mark.parametrize(
+    "metrics_fixture,test_property,expected_calculator",
+    [
+        # When no sample is available, returns None
+        (
+            "metrics_config_no_sample_with_target",
+            "obj_store_mem_max_pending_output_per_task",
+            lambda m: None,
+        ),
+        # When sample is available, uses average_bytes_per_output
+        (
+            "metrics_config_with_sample",
+            "obj_store_mem_max_pending_output_per_task",
+            lambda m: (
+                m.average_bytes_per_output
+                * m._op.data_context._max_num_blocks_in_streaming_gen_buffer
+            ),
+        ),
+        # When no sample is available, returns None
+        (
+            "metrics_config_pending_outputs_no_sample",
+            "obj_store_mem_pending_task_outputs",
+            lambda m: None,
+        ),
+    ],
+)
+def test_obj_store_mem_estimation(
+    request, metrics_fixture, test_property, expected_calculator
+):
+    """Test object store memory estimation for various scenarios."""
+    metrics = request.getfixturevalue(metrics_fixture)
+    actual = getattr(metrics, test_property)
+    expected = expected_calculator(metrics)
+
+    assert (
+        actual == expected
+    ), f"Expected {test_property} to be {expected}, got {actual}"
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))