Repository: ray-project/ray
File: python/ray/data/tests/test_streaming_executor.py (Python)

Code before:

        # Create a streaming generator that produces a single 128 MiB output block.
        streaming_gen = create_stub_streaming_gen(block_nbytes=[128 * MiB])
        data_op_task = DataOpTask(0, streaming_gen)

        # Wait for the block to be ready, then remove the worker node.
        ray.wait([streaming_gen], fetch_local=False)
        cluster.remove_node(worker_node)

        # The block shouldn't be available anymore, so we shouldn't read any data.
        bytes_read = data_op_task.on_data_ready(None)
        assert bytes_read == 0

        # Re-add the worker node, and run the task to completion.
        new_worker_node = cluster.add_node(num_cpus=1)  # noqa: F841
        cluster.wait_for_nodes()
        while not data_op_task.has_finished:
            ray.wait([streaming_gen], fetch_local=False)
            bytes_read += data_op_task.on_data_ready(None)

        # We should now be able to read the 128 MiB block.
        assert bytes_read == pytest.approx(128 * MiB)

    @patch("ray.data._internal.execution.interfaces.physical_operator.logger.warning")
    @patch("ray.data._internal.execution.interfaces.physical_operator.ray.get")
    def test_operator_name_in_metadata_timeout_warning(
        self, mock_ray_get, mock_warning, ray_start_regular_shared
    ):
        """Test that the operator name appears in the metadata timeout warning."""
        mock_ray_get.side_effect = ray.exceptions.GetTimeoutError()

        streaming_gen = create_stub_streaming_gen(block_nbytes=[1])
        task = DataOpTask(0, streaming_gen, operator_name="MapBatches(fn)")

        # Pre-set both refs to non-nil to skip the generator state machine
        # and jump straight to the ray.get(meta_ref) timeout branch.
        task._pending_block_ref = ray.put("dummy_block")
        task._pending_meta_ref = ray.put("dummy_meta")

        task.on_data_ready(None)

        mock_warning.assert_called_once()
        msg = mock_warning.call_args[0][0]
        assert "operator 'MapBatches(fn)'" in msg

    @patch("time.perf_counter")
    def test_on_data_ready_output_backpressure_tracking(
        self, mock_perf_counter, ray_start_regular_shared
    ):
        """Test that DataOpTask tracks output backpressure time correctly.


Review comment (suggestion):
While I generally think we should expect tests with all contributions, I'm wondering if this test (checking that a warning message contains a particular string) is worth the maintenance cost. The logic is kind of non-trivial

Code after:

        # Create a streaming generator that produces a single 128 MiB output block.
        streaming_gen = create_stub_streaming_gen(block_nbytes=[128 * MiB])
        data_op_task = DataOpTask(0, streaming_gen)

        # Wait for the block to be ready, then remove the worker node.
        ray.wait([streaming_gen], fetch_local=False)
        cluster.remove_node(worker_node)

        # The block shouldn't be available anymore, so we shouldn't read any data.
        bytes_read = data_op_task.on_data_ready(None)
        assert bytes_read == 0

        # Re-add the worker node, and run the task to completion.
        new_worker_node = cluster.add_node(num_cpus=1)  # noqa: F841
        cluster.wait_for_nodes()
        while not data_op_task.has_finished:
            ray.wait([streaming_gen], fetch_local=False)
            bytes_read += data_op_task.on_data_ready(None)

        # We should now be able to read the 128 MiB block.
        assert bytes_read == pytest.approx(128 * MiB)

    @patch("time.perf_counter")
    def test_on_data_ready_output_backpressure_tracking(
        self, mock_perf_counter, ray_start_regular_shared
    ):
        """Test that DataOpTask tracks output backpressure time correctly.

        Output backpressure occurs when max_bytes_to_read=0, meaning the
        downstream can't consume more data. The task should track the total
        time spent in this state and report it via TaskExecDriverStats.
        """
        # Use 2 blocks so the task stays alive across both BP periods.
        streaming_gen = create_stub_streaming_gen(block_nbytes=[128 * MiB, 128 * MiB])

        captured_stats = {}

        def capture_done(exc, task_exec_stats, task_exec_driver_stats):
            captured_stats["exc"] = exc
            captured_stats["task_exec_stats"] = task_exec_stats
            captured_stats["task_exec_driver_stats"] = task_exec_driver_stats

        data_op_task = DataOpTask(
            0,
            streaming_gen,
            task_done_callback=capture_done,
        )

        clock = 0.0
        mock_perf_counter.return_value = clock