ray-project--ray/python/ray/data/tests/test_op_runtime_metrics.py

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__]))