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

Code before:
            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.

        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,

Review comment (bug):
This can fail if we change the way we call logger, for example changing to `logger.warning("operator '%s'...", operator_name)`, then this test will fail as the operator name does not exist in the "operator '%s'..."

Maybe we can use `caplog` and do the following, it will record the log message we actually printed

```py
  with caplog.at_level(logging.WARNING, logger="ray.data._internal.execution.interfaces.physical_operator"):
      task.on_data_ready(None)

  assert any("operator 'MapBatches(fn)'" in r.message for r in caplog.records)

```

Code after:
            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

        # Wait for data to become available
        ray.wait([streaming_gen], fetch_local=False)

        # 1st backpressure period: 2.5s
        clock = 1.0
        mock_perf_counter.return_value = clock
        assert data_op_task.on_data_ready(0) == 0

        clock = 3.5
        mock_perf_counter.return_value = clock

        # Resume: ends 1st BP period (2.5s), reads block 1 (limited to 1 byte
        # so it reads exactly one block and stops)
        data_op_task.on_data_ready(None)
        assert not data_op_task.has_finished